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)