def main():
transform = transforms.Compose([
transforms.Resize((32, 32)), # 因为下载的图片大小不一，所以先resize
transforms.ToTensor(), # ToTensor() 将 H*W*C的numpy 变为 C*H*W的tensor
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]) # 标准化处理
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
net = LeNet() # 实例化对象
net.load_state_dict(torch.load('Lenet.pth')) # 通过load_state_dict载入保存的权重文件
#im = Image.open('1.jpg') # PIL库的Image 打开图像，H*W*C
im = Image.open('./pytorch_classification/Test1_official_demo/1.jpg')
im = transform(
im
) # [C, H, W] # 如果要在网络中传播，必须变为tensor格式，所以transform变为C*H*W
im = torch.unsqueeze(
im, dim=0
) # [N, C, H, W] # ^ torch.unsqueeze(~~~ , dim=0)在第0维增加新的维度；变为N*C*H*W
with torch.no_grad(): # ^ 测试阶段，不需要求损失梯度，使用torch.no_grad()
outputs = net(im) # 图像传入网络中得到输出，输出的维度为[batch, 10]，只关注dim=1
predict = torch.max(outputs, dim=1)[1].data.numpy(
) # 找到第1个维度中最大的数，但是只关注其位置（即[1]，index），转化为numpy
predict1 = torch.softmax(outputs, dim=1) # 得到了概率分布
print(classes[int(predict)]) # 将index传入到classes就得到了类别
print(predict1) # 打印概率分布
def main():
transform = transforms.Compose([
transforms.Resize(
(32, 32)), #transforms.Resize((32,32))因为下载的图不一定是标准的，先转为32,32
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
net = LeNet() #实例化
net.load_state_dict(torch.load('Lenet.pth')) #载入权重文件
im = Image.open(
'luxing_cat.jpg') #通过PIL、numpy一般导入的格式为（height,width,channel）[H,W,C]
im = transform(im) # [C, H, W]
im = torch.unsqueeze(im, dim=0) # 加上batch维度（dim=0表示在最前面） [N, C, H, W]
with torch.no_grad():
outputs = net(im)
predict = torch.max(outputs, dim=1)[1].data.numpy()
print(classes[int(predict)])
#使用softmax
with torch.no_grad():
outputs = net(im)
predict = torch.softmax(outputs, dim=1) #dim=0是batch维度
print(predict)
def test(device, test_dataset_path, model_path):
test_loader = get_test_loader(test_dataset_path)
loss_criterion = torch.nn.NLLLoss()
model = LeNet(35)
if device.type == 'cpu':
model.load_state_dict(torch.load(model_path, map_location=device))
else:
model.load_state_dict(torch.load(model_path))
model.to(device)
test_model(model, test_loader, loss_criterion, [], device, None)
def Inference(index):
model = LeNet()
model.load_state_dict(torch.load('MNIST_Model.pth'))
img, label = testData.__getitem__(index)
img = img.unsqueeze(0)
output = model(img)
print(output)
output = output.tolist()
output = [i for i in output[0]]
x = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
plt.figure()
plt.subplot(2, 1, 1)
plt.imshow(img.reshape(28, 28))
plt.subplot(2, 1, 2)
plt.ylim(0, 1)
plt.yticks([0, 0.2, 0.4, 0.6, 0.8, 1])
plt.xlim(0, 9)
plt.bar(x, output)
plt.show()
def main():
transform = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
net = LeNet()
net.load_state_dict(torch.load('Lenet.pth'))
im = Image.open('1.jpg')
im = transform(im) # [C, H, W]
im = torch.unsqueeze(im, dim=0) # [N, C, H, W]
with torch.no_grad():
outputs = net(im)
predict = torch.max(outputs, dim=1)[1].data.numpy()
print(classes[int(predict)])
import torch
import torchvision.transforms as transforms
from PIL import Image
from model import LeNet
transform = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
net = LeNet()
net.load_state_dict(torch.load('Lenet.pth'))
im = Image.open('1.jpg')
im = transform(im) # [C, H, W]
im = torch.unsqueeze(im, dim=0) # [N, C, H, W]
with torch.no_grad():
outputs = net(im)
predict = torch.max(outputs, dim=1)[1].data.numpy()
print(classes[int(predict)])
def main(args):
check_path(args)
# CIFAR-10的全部类别，一共10类
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
# 数据集
data_builder = DataBuilder(args)
dataSet = DataSet(data_builder.train_builder(),
data_builder.test_builder(), classes)
# 选择模型
if args.lenet:
net = LeNet()
model_name = args.name_le
elif args.vgg:
net = Vgg16_Net()
model_name = args.name_vgg
else:
raise "Sorry, you can only select LeNet or VGG."
# 交叉熵损失函数
criterion = nn.CrossEntropyLoss()
# SGD优化器
optimizer = optim.SGD(net.parameters(),
lr=args.learning_rate,
momentum=args.sgd_momentum,
weight_decay=args.weight_decay)
# 模型的参数保存路径，默认为 "./model/state_dict"
model_path = os.path.join(args.model_path, model_name)
# 指定在GPU / CPU上运行程序
device = t.device("cuda:0" if (
t.cuda.is_available() and not args.no_cuda) else "cpu")
# 启动训练
if args.do_train:
print("Training...")
trainer = Trainer(net, criterion, optimizer, dataSet.train_loader,
args)
trainer.train(epochs=args.epoch)
# 只保存模型参数
t.save(net.state_dict(), model_path)
# 启动测试
if args.do_eval:
if not os.path.exists(model_path):
print(
"Sorry, there's no saved model yet, you need to train first.")
return
print("Testing...")
device = t.device("cuda:0" if t.cuda.is_available() else "cpu")
net.load_state_dict(t.load(model_path, map_location=device))
# net.eval()
tester = Tester(dataSet.test_loader, net, args)
tester.test()
if args.show_model:
if not os.path.exists(model_path):
print(
"Sorry, there's no saved model yet, you need to train first.")
return
show_model(args)
if args.do_predict:
net.load_state_dict(t.load(model_path, map_location=device))
predictor = Predictor(net, classes)
# img_path = 'test'
# img_name = [os.path.join(img_path, x) for x in os.listdir(img_path)]
# for img in img_name:
# predictor.predict(img)
img_path = 'test/cat0.jpg'
predictor.predict(img_path)
loader = torch.utils.data.DataLoader(subset,
batch_size=args.batch_size,
shuffle=False,
num_workers=4,
pin_memory=has_cuda)
loaderOneByOne = torch.utils.data.DataLoader(subset,
batch_size=1,
shuffle=False,
num_workers=4,
pin_memory=has_cuda)
# Retrieve pre trained model
classes = 10
model = LeNet()
d = torch.load('model/best.pth.tar', map_location='cpu')
model.load_state_dict(d['state_dict'], strict=False)
model.eval()
for p in model.parameters():
p.requires_grad_(False)
#change criterion for ImageNet-1k and CIFAR100 to Top5Criterion
criterion = lambda x, y: Top1Criterion(x, y, model)
if args.norm == 'L2':
norm = 2
elif args.norm == 'Linf':
norm = np.inf
elif args.norm == 'L0':
norm = 0
elif args.norm == 'L1':
norm = 1
def img2MNIST(filename):
img = Image.open(filename).convert('L')
img = img.resize((28,28),Image.ANTIALIAS)
arr = []
for i in range(28):
for j in range(28):
pixel = float(img.getpixel((j, i))) / 255.0
arr.append(pixel)
arr1 = np.array(arr).reshape((1,1,28,28))
result = torch.as_tensor(arr1, dtype=torch.float32)
return result
parser = argparse.ArgumentParser()
parser.add_argument("--filename", type=str, default='./test_sample.bmp')
parser.add_argument("--model", type=str, default='./model_save/LeNet.pth')
arg = parser.parse_args()
if __name__ == "__main__":
print("Tested picture " + arg.filename)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_load = LeNet()
model_load.load_state_dict(torch.load(arg.model))
net = model_load.to(device)
net.eval()
image = img2MNIST(arg.filename)
output_test = net(image)
_, predicted = torch.max(output_test, 1)
print("The hand writing number is: " + str(predicted.item()))
best_valid_loss = np.inf
iteration = 0
epoch = 1
if not os.path.isfile(
'/content/drive/My Drive/PyTorch_Classifier/classifier.pt'):
#save model
model_save_name = 'classifier.pt'
path = F"/content/drive/My Drive/PyTorch_Classifier/{model_save_name}"
torch.save(model.state_dict(), path)
#load model
model_save_name = 'classifier.pt'
path = F"/content/drive/My Drive/PyTorch_Classifier/{model_save_name}"
model.load_state_dict(torch.load(path))
# trainint with early stopping
while (epoch < args.epochs + 1) and (iteration < args.patience):
train(train_loader, model, optimizer, epoch, args.cuda, args.log_interval)
valid_loss = test(valid_loader, model, args.cuda)
if valid_loss > best_valid_loss:
iteration += 1
print('Loss was not improved, iteration {0}'.format(str(iteration)))
else:
print('Saving model...')
iteration = 0
best_valid_loss = valid_loss
state = {
'net': model.module if args.cuda else model,
'acc': valid_loss,
# select your model
model_select = side_bar.radio('Model Name', ('LeNet-5', 'VGG-16'))
# select your testing device
device_select = side_bar.radio('Device', ('CPU', 'GPU'))
upload_img = st.file_uploader("Please upload your image", type="jpg")
if upload_img is not None:
# 展示图片
st.image(upload_img)
# 选择模型
if model_select == 'VGG-16':
net = Vgg16_Net()
model_path = '../model/state_dict_vgg'
else:
net = LeNet()
model_path = '../model/state_dict_le'
try:
net.load_state_dict(
t.load(model_path, map_location=t.device(device_select.lower())))
predictor = Predictor(net, classes)
result = predictor.predict(upload_img)
st.write("The result is", classes[result])
except RuntimeError:
st.write(
"Please check your device, do not select GPU button when you have no CUDA!"
)
import torch.nn as nn
from model import LeNet
# ... 此处略去定义测试数据载入器的代码，具体参考代码4.3
# save_info = { # 保存的信息
# "iter_num": iter_num, # 迭代步数
# "optimizer": optimizer.state_dict(), # 优化器的状态字典
# "model": model.state_dict(), # 模型的状态字典
# }
model_path = "./model.pth" # 假设模型保存在model.pth文件中
save_info = torch.load(model_path) # 载入模型
model = LeNet() # 定义LeNet模型
criterion = nn.CrossEntropyLoss() # 定义损失函数
model.load_state_dict(save_info["model"]) # 载入模型参数
model.eval() # 切换模型到测试状态
test_loss = 0
correct = 0
total = 0
with torch.no_grad(): # 关闭计算图
for batch_idx, (inputs, targets) in enumerate(data_test_loader):
outputs = model(inputs)
loss = criterion(outputs, targets)
test_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
def train_lre(hyperparameters,
data_loader,
val_data,
val_labels,
test_loader,
n_val=5):
net, opt = build_model(hyperparameters)
plot_step = 10
accuracy_log = dict()
data_log = dict()
subselect_val = n_val * val_labels.unique().size()[0] != len(val_labels)
meta_net = LeNet(n_out=1)
for i in range(hyperparameters['num_iterations']):
t = time.time()
net.train()
# Line 2 get batch of data
image, labels, index = next(iter(data_loader))
# since validation data is small I just fixed them instead of building an iterator
# initialize a dummy network for the meta learning of the weights
meta_net.load_state_dict(net.state_dict())
if torch.cuda.is_available():
meta_net.cuda()
image = to_var(image, requires_grad=False)
labels = to_var(labels, requires_grad=False)
# Lines 4 - 5 initial forward pass to compute the initial weighted loss
y_f_hat = meta_net(image) # Line 4
cost = functional.binary_cross_entropy_with_logits(
y_f_hat, labels, reduce=False) # Line 5
eps = to_var(torch.zeros(cost.size())) # Line 5
l_f_meta = torch.sum(cost * eps)
meta_net.zero_grad()
# Line 6 perform a parameter update
grads = torch.autograd.grad(l_f_meta, (meta_net.params()),
create_graph=True)
meta_net.update_params(hyperparameters['lr'], source_params=grads)
# Line 8 - 10 2nd forward pass and getting the gradients with respect to epsilon
if subselect_val:
class_inds = list()
for c in val_labels.unique():
matching_inds = (val_labels == c).nonzero()
class_inds.append(matching_inds[np.random.permutation(
len(matching_inds))[:n_val]])
class_inds = torch.cat(class_inds)
val_input = val_data[class_inds].squeeze_(1)
val_output = val_labels[class_inds].squeeze_(1)
else:
val_input = val_data
val_output = val_labels
y_g_hat = meta_net(val_input)
l_g_meta = functional.binary_cross_entropy_with_logits(
y_g_hat, val_output)
grad_eps = torch.autograd.grad(l_g_meta, eps, only_inputs=True)[0]
# Line 11 computing and normalizing the weights
w_tilde = torch.clamp(-grad_eps, min=0)
norm_c = torch.sum(w_tilde)
if norm_c != 0:
w = w_tilde / norm_c
else:
w = w_tilde
# Lines 12 - 14 computing for the loss with the computed weights
# and then perform a parameter update
y_f_hat = net(image)
cost = functional.binary_cross_entropy_with_logits(y_f_hat,
labels,
reduce=False)
data_log[i] = pd.DataFrame(
data={
'w': w.numpy(),
'grad_eps': grad_eps.numpy(),
'pre_w_cost': cost.detach().numpy(),
'label': labels.numpy(),
'index': index.numpy(),
})
l_f = torch.sum(cost * w)
# print(data_log[i].head())
# print(data_log[i].tail())
opt.zero_grad()
l_f.backward()
opt.step()
# meta_l = smoothing_alpha * meta_l + (1 - smoothing_alpha) * l_g_meta.item()
# meta_losses_clean.append(meta_l / (1 - smoothing_alpha ** (i + 1)))
#
# net_l = smoothing_alpha * net_l + (1 - smoothing_alpha) * l_f.item()
# net_losses.append(net_l / (1 - smoothing_alpha ** (i + 1)))
#
print(time.time() - t)
if i % plot_step == 0:
print(i)
net.eval()
acc_df = {'preds': [], 'labels': [], 'index': []}
for itr, (test_img, test_label,
test_idx) in enumerate(test_loader):
test_img = to_var(test_img, requires_grad=False)
test_label = to_var(test_label, requires_grad=False)
output = net(test_img)
predicted = functional.sigmoid(output)
acc_df['preds'].extend(predicted.detach().numpy().tolist())
acc_df['labels'].extend(test_label.numpy().tolist())
acc_df['index'].extend(test_idx.numpy().tolist())
accuracy_log[i] = pd.DataFrame(acc_df).sort_values(
['labels', 'index']).set_index('index', drop=True)
data_log = pd.concat(data_log)
data_log.index.set_names('iteration', level=0, inplace=True)
accuracy_log = pd.concat(accuracy_log)
accuracy_log.index.set_names('iteration', level=0, inplace=True)
return data_log, accuracy_log
import base64
import flask
import io
from model import LeNet
from PIL import Image, ImageFilter
import re
import torch
from torchvision import transforms
# Create Flask App
app = flask.Flask(__name__)
# Load prediction model
device = torch.device('cpu')
model = LeNet()
model.load_state_dict(torch.load("model.pt", map_location=device))
model.eval()
@app.route("/")
def index():
return flask.render_template("index.html")
@app.route("/predict", methods=["POST"])
def predict():
request = flask.request.get_json()
# Get image from request
image = request["image"]
# Remove metadata prefix
image = re.sub('^data:image/.+;base64,', '', image)
dirs = os.listdir("./gcommands/train/")
dirs = sorted(dirs)
parser = argparse.ArgumentParser(
description='ConvNets for Speech Commands Recognition')
parser.add_argument('--wav_path',
default='gcommands/recordings/one/3.wav',
help='path to the audio file')
args = parser.parse_args()
path = args.wav_path
warnings.filterwarnings("ignore")
model = LeNet()
model.load_state_dict(torch.load("checkpoint/ckpt.t7"))
model.eval()
wav = spect_loader(path,
window_size=.02,
window_stride=.01,
normalize=True,
max_len=101,
window='hamming')
#print(wav.shape)
with torch.no_grad():
label = model.forward(wav.view(1, 1, 161, 101))
print(dirs[np.argmax(np.ravel(label.numpy()))])
if args.optimizer.lower() == 'adam':
optimizer = optim.Adam(model.parameters(), lr=args.lr)
elif args.optimizer.lower() == 'sgd':
optimizer = optim.SGD(model.parameters(), lr=args.lr,
momentum=args.momentum)
else:
optimizer = optim.SGD(model.parameters(), lr=args.lr,
momentum=args.momentum)
best_valid_loss = np.inf
iteration = 0
epoch = 1
if args.test_mode:
checkpoint = torch.load(args.chkpt_path)
model.load_state_dict(checkpoint['net'].state_dict())
# optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
model.eval()
# attack model
attack(test_loader, 0, model)
# trainint with early stopping
while (epoch < args.epochs + 1) and (iteration < args.patience) and not args.test_mode:
train(train_loader, model, optimizer, epoch, args.cuda, args.log_interval)
# attack model
attack(test_loader, 0, model)
valid_loss = test(valid_loader, model, args.cuda)
if valid_loss > best_valid_loss:
iteration += 1
print('Loss was not improved, iteration {0}'.format(str(iteration)))
[transforms.Resize((32, 32)),
transforms.ToTensor()]),
download=True) #从internet上下载MNIST测试数据集，并Resize成32x32大小，转为Tensor
# data_train_loader = DataLoader(data_train, batch_size=256, shuffle=True, num_workers=8) #8个线程处理，随机抽取，batch=256
data_test_loader = DataLoader(data_test,
batch_size=args.batch_size,
num_workers=0)
### LeNet工程推理代码部分
save_info_path = 'D:/数据结构学习/LeNet_master/model_save/model.pth'
save_info = torch.load(save_info_path)
model = LeNet()
criterion = nn.CrossEntropyLoss()
model.load_state_dict(save_info["model"])
model.eval()
test_loss = 0
correct = 0
total = 0
epoch_num = 5
outputs_plt = []
inputs_plt = []
target_plt = []
with torch.no_grad(): ## 测试时不需要开启计算图
for epoch in range(epoch_num):
for batch_idx, (inputs, targets) in enumerate(data_test_loader):
# print(inputs.shape)
# print(targets.shape)
'net': model.module.state_dict() if args.cuda else model.state_dict(),
'acc': valid_loss,
'epoch': epoch,
}
if not os.path.isdir('checkpoint'):
os.mkdir('checkpoint')
torch.save(state, './checkpoint/{}.t7'.format(args.arc))
epoch += 1
# test model
checkpoint = torch.load('./checkpoint/{}.t7'.format(args.arc),map_location = lambda storage, loc: storage)
if args.arc == 'LeNet':
test_model = LeNet(num_classes=num_classes)
elif args.arc.startswith('VGG'):
test_model = VGG(args.arc, num_classes=num_classes)
else:
test_model = LeNet(num_classes=num_classes)
test_model.load_state_dict(checkpoint["net"])
if args.cuda:
print('Using CUDA with {0} GPUs'.format(torch.cuda.device_count()))
test_model = torch.nn.DataParallel(test_model).cuda()
test(valid_loader, test_model, args.cuda)
test(test_loader, test_model, args.cuda)
def main():
# Training settings
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument('--dataset',
type=str,
default="mnist",
choices=["mnist", "cifar10"],
metavar='D',
help='training dataset (mnist or cifar10)')
parser.add_argument('--batch-size',
type=int,
default=64,
metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('--test-batch-size',
type=int,
default=1000,
metavar='N',
help='input batch size for testing (default: 1000)')
parser.add_argument('--percent',
type=list,
default=[0.8, 0.92, 0.991, 0.93],
metavar='P',
help='pruning percentage (default: 0.8)')
parser.add_argument('--alpha',
type=float,
default=5e-4,
metavar='L',
help='l2 norm weight (default: 5e-4)')
parser.add_argument('--rho',
type=float,
default=1e-2,
metavar='R',
help='cardinality weight (default: 1e-2)')
parser.add_argument(
'--l1',
default=False,
action='store_true',
help='prune weights with l1 regularization instead of cardinality')
parser.add_argument('--l2',
default=False,
action='store_true',
help='apply l2 regularization')
parser.add_argument('--num_pre_epochs',
type=int,
default=3,
metavar='P',
help='number of epochs to pretrain (default: 3)')
parser.add_argument('--num_epochs',
type=int,
default=10,
metavar='N',
help='number of epochs to train (default: 10)')
parser.add_argument('--num_re_epochs',
type=int,
default=3,
metavar='R',
help='number of epochs to retrain (default: 3)')
parser.add_argument('--lr',
type=float,
default=1e-3,
metavar='LR',
help='learning rate (default: 1e-2)')
parser.add_argument('--adam_epsilon',
type=float,
default=1e-8,
metavar='E',
help='adam epsilon (default: 1e-8)')
parser.add_argument('--no-cuda',
action='store_true',
default=False,
help='disables CUDA training')
parser.add_argument('--seed',
type=int,
default=1,
metavar='S',
help='random seed (default: 1)')
parser.add_argument('--save-model',
action='store_true',
default=False,
help='For Saving the current Model')
parser.add_argument('--structured',
action='store_true',
default=False,
help='Enabling Structured Pruning')
parser.add_argument('--test',
action='store_true',
default=False,
help='For Testing the current Model')
parser.add_argument(
'--stat',
action='store_true',
default=False,
help='For showing the statistic result of the current Model')
parser.add_argument('--n1',
type=int,
default=2,
metavar='N',
help='ReRAM OU size (row number) (default: 2)')
parser.add_argument('--n2',
type=int,
default=2,
metavar='N',
help='ReRAM OU size (column number) (default: 2)')
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
if args.dataset == "mnist":
train_loader = torch.utils.data.DataLoader(datasets.MNIST(
'data',
train=True,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
])),
batch_size=args.batch_size,
shuffle=True,
**kwargs)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST('data',
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
])),
batch_size=args.test_batch_size,
shuffle=True,
**kwargs)
else:
args.percent = [0.8, 0.92, 0.93, 0.94, 0.95, 0.99, 0.99, 0.93]
args.num_pre_epochs = 5
args.num_epochs = 20
args.num_re_epochs = 5
train_loader = torch.utils.data.DataLoader(datasets.CIFAR10(
'data',
train=True,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.49139968, 0.48215827, 0.44653124),
(0.24703233, 0.24348505, 0.26158768))
])),
shuffle=True,
batch_size=args.batch_size,
**kwargs)
test_loader = torch.utils.data.DataLoader(
datasets.CIFAR10('data',
train=False,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
(0.49139968, 0.48215827, 0.44653124),
(0.24703233, 0.24348505, 0.26158768))
])),
shuffle=True,
batch_size=args.test_batch_size,
**kwargs)
model = LeNet().to(device) if args.dataset == "mnist" else AlexNet().to(
device)
optimizer = PruneAdam(model.named_parameters(),
lr=args.lr,
eps=args.adam_epsilon)
structured_tag = "_structured{}x{}".format(
args.n1, args.n2) if args.structured else ""
model_file = "mnist_cnn{}.pt".format(structured_tag) if args.dataset == "mnist" \
else 'cifar10_cnn{}.pt'.format(structured_tag)
if args.stat or args.test:
print("=> loading model '{}'".format(model_file))
if os.path.isfile(model_file):
model.load_state_dict(torch.load(model_file))
print("=> loaded model '{}'".format(model_file))
if args.test:
test(args, model, device, test_loader)
if args.stat:
show_statistic_result(args, model)
else:
print("=> loading model failed '{}'".format(model_file))
else:
checkpoint_file = 'checkpoint{}.pth.tar'.format(
"_mnist" if args.dataset == "mnist" else "_cifar10")
if not os.path.isfile(checkpoint_file):
pre_train(args, model, device, train_loader, test_loader,
optimizer)
torch.save(
{
'epoch': args.num_pre_epochs,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, checkpoint_file)
else:
print("=> loading checkpoint '{}'".format(checkpoint_file))
checkpoint = torch.load(checkpoint_file)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}'".format(checkpoint_file))
train(args, model, device, train_loader, test_loader, optimizer)
mask = apply_l1_prune(model, device, args) if args.l1 else apply_prune(
model, device, args)
print_prune(model)
test(args, model, device, test_loader)
retrain(args, model, mask, device, train_loader, test_loader,
optimizer)
if args.save_model:
torch.save(model.state_dict(), model_file)
def main():
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
data_transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# load image
img_path = 'a.jpg'
assert os.path.exists(img_path), "file: '{}' does not exit.".format(
img_path)
img = Image.open(img_path)
plt.imshow(img)
# [N, C, H, W]
img = data_transform(img)
# expand batch dimension
img = torch.unsqueeze(img, dim=0)
print(img.shape)
# read class_indict
json_path = 'class_indices.json'
assert os.path.exists(json_path), "file: '{}' does not exist.".format(
json_path)
json_file = open(json_path, 'r')
class_indict = json.load(json_file)
# print(class_indict)
# create model
model = LeNet(num_classes=5).to(device)
# load model weights
model_weight_path = 'weights/lenet.pth'
model.load_state_dict(torch.load(model_weight_path, map_location=device))
model.eval()
print('=================================')
dummy_input = torch.randn(1, 3, 224, 224).to(device)
torch.onnx.export(model,
dummy_input,
'lenet.onnx',
dynamic_axes={
'image': {
0: 'B'
},
'outputs': {
0: 'B'
}
},
input_names=['image'],
output_names=['outputs'],
opset_version=12)
print('=================================')
print('---------------------------------')
traced_script_module = torch.jit.trace(model, dummy_input)
traced_script_module.save("lnet.pt")
print('---------------------------------')
with torch.no_grad():
# predict class
import time
t1 = time.time()
model(img.to(device)).cpu()
t2 = time.time()
print('torch 推理花费{}ms'.format(1000 * (t2 - t1)))
output = torch.squeeze(model(img.to(device))).cpu()
predict = torch.softmax(output, dim=0)
predict_cla = torch.argmax(predict).numpy()
print_res = "class: {} prob: {:.3}".format(class_indict[str(predict_cla)],
predict[predict_cla].numpy())
plt.title(print_res)
print(print_res)
plt.show()
cudnn.benchmark = True
model = torch.nn.DataParallel(base_model)
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=5e-4)
scheduler = MultiStepLR(optimizer, milestones=[200, 400], gamma=0.1)
# Resume from checkpoint if required
start_epoch = 0
if checkpoint is not None:
print('==> Resuming from checkpoint..')
print(checkpoint)
checkpoint = torch.load(checkpoint)
base_model.load_state_dict(checkpoint['net'])
start_epoch = checkpoint['epoch']
scheduler.step(start_epoch)
# Main routine
if args.task == 'train':
# Training routine
for epoch in range(start_epoch + 1, args.epochs + 1):
print('===train(epoch={})==='.format(epoch))
t1 = time.time()
scheduler.step()
model.train()
macer_train(args.sigma, args.lbd, args.gauss_num, args.beta,
args.gamma, num_classes, model, trainloader, optimizer,
device)
