def pruning_process(model): print("------------------------- Before pruning --------------------------------") criterion = nn.CrossEntropyLoss().cuda() accuracy = util.validate(args, test_loader, model, criterion) print("------------------------- pruning --------------------------------------") if args.model_mode == 'c' or args.model_mode =='a': model.prune_by_percentile( ['conv1'], q=100-58.0) model.prune_by_percentile( ['conv2'], q=100-22.0) model.prune_by_percentile( ['conv3'], q=100-34.0) model.prune_by_percentile( ['conv4'], q=100-36.0) model.prune_by_percentile( ['conv5'], q=100-53.0) model.prune_by_percentile( ['conv6'], q=100-24.0) model.prune_by_percentile( ['conv7'], q=100-42.0) model.prune_by_percentile( ['conv8'], q=100-32.0) model.prune_by_percentile( ['conv9'], q=100-27.0) model.prune_by_percentile( ['conv10'], q=100-34.0) model.prune_by_percentile( ['conv11'], q=100-35.0) model.prune_by_percentile( ['conv12'], q=100-29.0) model.prune_by_percentile( ['conv13'], q=100-36.0) if args.model_mode == 'd' or args.model_mode == 'a': model.prune_by_percentile( ['fc1'], q=100-10.0) model.prune_by_percentile( ['fc2'], q=100-10.0) model.prune_by_percentile( ['fc3'], q=100-10.0) print("------------------------- After pruning --------------------------------") util.print_nonzeros(model, f"{args.save_dir}/{args.log}") accuracy = util.validate(args, test_loader, model, criterion) torch.save(model, f"{args.save_dir}/model_pruned.ptmodel") util.log(f"{args.save_dir}/{args.log}", f"weight\t{args.save_dir}/{args.out_pruned_folder}") util.log(f"{args.save_dir}/{args.log}", f"model\t{args.save_dir}/model_pruned.ptmodel") util.log(f"{args.save_dir}/{args.log}", f"accuracy after pruning\t{accuracy}") util.layer2torch(model, f"{args.save_dir}/{args.out_pruned_folder}") weight_list = util.parameters2list(model.children()) util.save_parameters(f"{args.save_dir}/{args.out_pruned_folder}", weight_list) # Retrain print("------------------------- Prune and Retrain ----------------------------") tok="prune_re" util.initial_train(model, args, train_loader, test_loader, tok, use_cuda=True) print("------------------------- After Retraining -----------------------------") util.print_nonzeros(model, f"{args.save_dir}/{args.log}") accuracy = util.validate(args, test_loader, model, criterion) torch.save(model, f"{args.save_dir}/model_prune_retrain_{args.reepochs}.ptmodel") util.log(f"{args.save_dir}/{args.log}", f"weight\t{args.save_dir}/{args.out_pruned_re_folder}") util.log(f"{args.save_dir}/{args.log}", f"model\t{args.save_dir}/mmodel_prune_retrain_{args.reepochs}.ptmodel") util.log(f"{args.save_dir}/{args.log}", f"accuracy after prune retrain\t{accuracy}") util.layer2torch(model, f"{args.save_dir}/{args.out_pruned_re_folder}") weight_list = util.parameters2list(model.children()) util.save_parameters(f"{args.save_dir}/{args.out_pruned_re_folder}", weight_list) return model
def encoding_process(): util.topic_log("Accuracy before huffman encoding") util.print_nonzeros(model, args.log_file_path) top1_acc, top5_acc = util.val_epoch(val_loader, model, args, topk=(1, 5)) util.log(args.log_file_path, f"accuracy before huffman encoding\t{top1_acc} ({top5_acc})") util.topic_log("encoding") mesa2_huffman_encode_model(model, args)
def pruning_process(): util.topic_log("Before pruning") top1_acc, top5_acc = util.val_epoch(val_loader, model, args, topk=(1, 5)) util.log(args.log_file_path, f"before pruning accuracy\t{top1_acc} ({top5_acc})") util.topic_log("Pruning CNN") model.prune(args) util.topic_log("After prune CNN") util.print_nonzeros(model, args.log_file_path) top1_acc, top5_acc = util.val_epoch(val_loader, model, args, topk=(1, 5)) util.log(args.log_file_path, f"after pruning accuracy\t{top1_acc} ({top5_acc})") util.topic_log("Start retrain after prune CNN") util.train(model, args, train_loader, val_loader, 'prune_retrain') util.topic_log("After retraining") top1_acc, top5_acc = util.val_epoch(val_loader, model, args, topk=(1, 5)) util.log(args.log_file_path, f"after pruning and retrain accuracy\t{top1_acc} ({top5_acc})")
def pruning_process(model): print( "------------------------- Before pruning --------------------------------" ) util.print_nonzeros(model, f"{args.save_dir}/{args.log}") accuracy = util.validate(val_loader, model, args) print( "------------------------- pruning CNN--------------------------------------" ) model.prune_by_percentile(['conv1'], q=100 - 66.0) model.prune_by_percentile(['conv2'], q=100 - 12.0) print( "------------------------------- After prune CNN ----------------------------" ) util.print_nonzeros(model, f"{args.save_dir}/{args.log}") prec1 = util.validate(val_loader, model, args) torch.save(model, f"{args.save_dir}/model_pruned.ptmodel") util.log(f"{args.save_dir}/{args.log}", f"weight\t{args.save_dir}/{args.out_pruned_folder}") util.log(f"{args.save_dir}/{args.log}", f"model\t{args.save_dir}/model_pruned.ptmodel") util.log(f"{args.save_dir}/{args.log}", f"prune acc\t{prec1}") util.layer2torch(f"{args.save_dir}/{args.out_pruned_folder}", model) weight_list = util.parameters2list(model.children()) util.save_parameters(f"{args.save_dir}/{args.out_pruned_folder}", weight_list) print( "------------------------- start retrain after prune CNN----------------------------" ) util.initial_train(model, args, train_loader, val_loader, 'prune_re') print( "------------------------- After Retraining -----------------------------" ) util.print_nonzeros(model, f"{args.save_dir}/{args.log}") accuracy = util.validate(val_loader, model, args) torch.save(model, f"{args.save_dir}/model_prune_retrain_{args.reepochs}.ptmodel") util.log(f"{args.save_dir}/{args.log}", f"weight\t{args.save_dir}/{args.out_pruned_re_folder}") util.log( f"{args.save_dir}/{args.log}", f"model\t{args.save_dir}/model_prune_retrain_{args.reepochs}.ptmodel") util.log(f"{args.save_dir}/{args.log}", f"prune and retrain acc\t{accuracy}") util.layer2torch(f"{args.save_dir}/{args.out_pruned_re_folder}", model) weight_list = util.parameters2list(model.children()) util.save_parameters(f"{args.save_dir}/{args.out_pruned_re_folder}", weight_list) return model
def pruning_process(model): print("------------------------- Before pruning --------------------------------") util.print_nonzeros(model, f"{args.save_dir}/{args.log}") accuracy = util.validate(val_loader, model, args) print("------------------------- pruning CNN--------------------------------------") model.prune_by_percentile( ['conv1'], q=100-58.0) model.prune_by_percentile( ['conv2'], q=100-22.0) model.prune_by_percentile( ['conv3'], q=100-34.0) model.prune_by_percentile( ['conv4'], q=100-36.0) model.prune_by_percentile( ['conv5'], q=100-53.0) model.prune_by_percentile( ['conv6'], q=100-24.0) model.prune_by_percentile( ['conv7'], q=100-42.0) model.prune_by_percentile( ['conv8'], q=100-32.0) model.prune_by_percentile( ['conv9'], q=100-27.0) model.prune_by_percentile( ['conv10'], q=100-34.0) model.prune_by_percentile( ['conv11'], q=100-35.0) model.prune_by_percentile( ['conv12'], q=100-29.0) model.prune_by_percentile( ['conv13'], q=100-36.0) print("------------------------------- After prune CNN ----------------------------") util.print_nonzeros(model, f"{args.save_dir}/{args.log}") prec1 = util.validate(val_loader, model, args) util.save_checkpoint({ 'state_dict': model.state_dict(), 'best_prec1': prec1, }, True, filename=os.path.join(args.save_dir, 'checkpoint_{}_alpha_{}.tar'.format('pruned',args.alpha))) util.log(f"{args.save_dir}/{args.log}", f"weight\t{args.save_dir}/{args.out_pruned_folder}") util.log(f"{args.save_dir}/{args.log}", f"model\t{args.save_dir}/model_pruned.ptmodel") util.log(f"{args.save_dir}/{args.log}", f"prune acc\t{prec1}") util.layer2torch(f"{args.save_dir}/{args.out_pruned_folder}" , model) weight_list = util.parameters2list(model.children()) util.save_parameters(f"{args.save_dir}/{args.out_pruned_folder}", weight_list) print("------------------------- start retrain after prune CNN----------------------------") util.initial_train(model, args, train_loader, val_loader, 'prune_re') print("------------------------- After Retraining -----------------------------") util.print_nonzeros(model, f"{args.save_dir}/{args.log}") accuracy = util.validate(val_loader, model, args) util.log(f"{args.save_dir}/{args.log}", f"weight\t{args.save_dir}/{args.out_pruned_re_folder}") util.log(f"{args.save_dir}/{args.log}", f"model\t{args.save_dir}/model_prune_retrain_{args.reepochs}.ptmodel") util.log(f"{args.save_dir}/{args.log}", f"prune and retrain acc\t{accuracy}") util.layer2torch(f"{args.save_dir}/{args.out_pruned_re_folder}" , model) weight_list = util.parameters2list(model.children()) util.save_parameters(f"{args.save_dir}/{args.out_pruned_re_folder}", weight_list) return model
test_loss /= len(test_loader.dataset) accuracy = 100. * correct / len(test_loader.dataset) print( f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)' ) return accuracy # Initial training print("--- Initial training ---") train(args.epochs) accuracy = test() util.log(args.log, f"initial_accuracy {accuracy}") torch.save(model, f"saves/initial_model.ptmodel") print("--- Before pruning ---") util.print_nonzeros(model) # Pruning model.prune_by_std(args.sensitivity) accuracy = test() util.log(args.log, f"accuracy_after_pruning {accuracy}") print("--- After pruning ---") util.print_nonzeros(model) # Retrain print("--- Retraining ---") optimizer.load_state_dict(initial_optimizer_state_dict) # Reset the optimizer train(args.epochs) torch.save(model, f"saves/model_after_retraining.ptmodel") accuracy = test() util.log(args.log, f"accuracy_after_retraining {accuracy}")
f = open("result_initial.csv", 'w') print(epoch_list_initial) for i in range(0, len(epoch_list_initial)): f.write( str(epoch_list_initial[i]) + ',' + str(psnr_list_initial[i]) + "\n") f.close() # Pruning srcnn.prune_by_std(opt.sensitivity) psnr = test() util.log(opt.log, "accuracy_after_pruning {}".format(psnr)) print("--- After pruning ---") util.print_nonzeros(srcnn) # Retrain print("--- Retraining ---") optimizer.load_state_dict(initial_optimizer_state_dict) # Reset the optimizer epoch_list_retrain = [] psnr_list_retrain = [] for epoch in range(1, opt.epochs + 1): train(epoch) psnr = test() epoch_list_retrain.append(epoch) psnr_list_retrain.append(psnr) if (epoch % 100 == 0): checkpoint(1, epoch)
model = LeNet_300_100(input_size, num_classes) opt = keras.optimizers.Adam(learning_rate=lr) model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy']) #path to save original model model_path = "lenet300-100.h5" model, history = train(model, x_train, y_train, batch_size, num_epochs, model_path) #training original_loss, original_accuracy = model.evaluate(x_test, y_test, verbose=0) #testing print("Result of Original Model:") print('Original Model loss:', original_loss) print('Original Model accuracy:', original_accuracy) print_nonzeros(model) get_plot(history) #Iterative training (prune+retrain) in single iteration and we run until we get pruning accuracy >= original model accuracy. #Threshold is calculated based on standard deviation of weights and sensitivity/quality factor. for i in range(1, 10): pretrained_model = keras.models.load_model(model_path) quality_parameter = 1.6 #sensitivity factor to calculate threshold #Pruning for layerid in range(len(pretrained_model.layers)): layer = pretrained_model.layers[layerid] weight = layer.get_weights() if len(weight) > 0: temp_weight = deepcopy(weight)
def main(): if args.seed is not None: random.seed(args.seed) torch.manual_seed(args.seed) cudnn.deterministic = True warnings.warn('You have chosen to seed training. ' 'This will turn on the CUDNN deterministic setting, ' 'which can slow down your training considerably! ' 'You may see unexpected behavior when restarting ' 'from checkpoints.') if args.gpu is not None: warnings.warn('You have chosen a specific GPU. This will completely ' 'disable data parallelism.') args.distributed = args.world_size > 1 if args.distributed: dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size) # create model if args.pretrained: print("=> using pre-trained model '{}'".format(args.arch)) model = models.__dict__[args.arch](pretrained=True) else: print("=> creating model '{}'".format(args.arch)) model = models.__dict__[args.arch]() if args.gpu is not None: model = model.cuda(args.gpu) elif args.distributed: model.cuda() model = torch.nn.parallel.DistributedDataParallel(model) else: if args.arch.startswith('alexnet') or args.arch.startswith('vgg'): model.features = torch.nn.DataParallel(model.features) model.cuda() else: model = torch.nn.DataParallel(model).cuda() # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda(args.gpu) optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) #if args.pretrained and 'alexnet' in args.arch: # model.load_state_dict(torch.load('saves/new_elt_0.0_0.pth')) #print(model) #util.print_model_parameters(model) # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True # Data loading code traindir = os.path.join(args.data, 'train') valdir = os.path.join(args.data, 'val') normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_dataset = datasets.ImageFolder( traindir, transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler( train_dataset) else: train_sampler = None train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), num_workers=args.workers, pin_memory=True, sampler=train_sampler) val_loader = torch.utils.data.DataLoader(datasets.ImageFolder( valdir, transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ])), batch_size=256, shuffle=False, num_workers=args.workers, pin_memory=True) if args.evaluate: util.print_nonzeros(model) validate(val_loader, model, criterion) return #if args.pretrained: # print('Pretrained model evaluation...') # validate(val_loader, model, criterion) curves = np.zeros(((args.epochs - args.start_epoch) * (len(train_loader) // args.print_freq), 4)) valid = np.zeros(((args.epochs - args.start_epoch), 3)) step = 0 for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) adjust_learning_rate(optimizer, epoch) # train for one epoch curves, step = train(train_loader, model, criterion, optimizer, epoch, args.reg, args.decay, curves, step) # evaluate on validation set valid[epoch, 0] = epoch valid[epoch, 1], valid[epoch, 2] = validate(val_loader, model, criterion) # plot training curve np.savetxt(os.path.join(save_path, 'curves.dat'), curves) clr1 = (0.5, 0., 0.) clr2 = (0.0, 0.5, 0.) fig, ax1 = plt.subplots() fig2, ax3 = plt.subplots() ax2 = ax1.twinx() ax4 = ax3.twinx() ax1.set_xlabel('steps') ax1.set_ylabel('Loss', color=clr1) ax1.tick_params(axis='y', colors=clr1) ax2.set_ylabel('Reg', color=clr2) ax2.tick_params(axis='y', colors=clr2) ax3.set_xlabel('steps') ax3.set_ylabel('Elt_sparsity', color=clr1) ax3.tick_params(axis='y', colors=clr1) start = 0 end = step markersize = 12 coef = 2. ax1.plot(curves[start:end, 0], curves[start:end, 1], '--', color=[c * coef for c in clr1], markersize=markersize) ax2.plot(curves[start:end, 0], curves[start:end, 2], '-', color=[c * coef for c in clr2], markersize=markersize) ax3.plot(curves[start:end, 0], curves[start:end, 3], '--', color=[c * coef for c in clr1], markersize=markersize) #ax2.set_ylim(bottom=20, top=100) ax1.legend(('Train loss'), loc='lower right') ax2.legend(('Train reg'), loc='lower left') fig.savefig(os.path.join(save_path, 'loss-vs-steps.pdf')) #ax4.set_ylim(bottom=20, top=100) ax3.legend(('Elt_sparsity'), loc='lower right') fig2.savefig(os.path.join(save_path, 'sparsity-vs-steps.pdf')) # plot validation curve np.savetxt(os.path.join(save_path, 'valid.dat'), valid) fig3, ax5 = plt.subplots() ax6 = ax5.twinx() ax5.set_xlabel('epochs') ax5.set_ylabel('Acc@1', color=clr1) ax5.tick_params(axis='y', colors=clr1) ax6.set_ylabel('Acc@5', color=clr2) ax6.tick_params(axis='y', colors=clr2) start = 0 end = epoch + 1 markersize = 12 coef = 2. ax5.plot(valid[start:end, 0], valid[start:end, 1], '--', color=[c * coef for c in clr1], markersize=markersize) ax6.plot(valid[start:end, 0], valid[start:end, 2], '-', color=[c * coef for c in clr2], markersize=markersize) #ax2.set_ylim(bottom=20, top=100) ax5.legend(('Acc@1'), loc='lower right') ax6.legend(('Acc@5'), loc='lower left') fig3.savefig(os.path.join(save_path, 'accuracy-vs-epochs.pdf')) torch.save( model.state_dict(), os.path.join( save_path, 'str_' + str(args.decay) + '_' + str(args.reg) + '.pth'))
def main(): global args, best_prec1 args = parser.parse_args() args.distributed = args.world_size > 1 if args.distributed: dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size) # create model if args.pretrained: print("=> using pre-trained model '{}'".format(args.arch)) model = models.__dict__[args.arch](pretrained=True) else: print("=> creating model '{}'".format(args.arch)) if args.arch.startswith('mobilenetv1'): model = MobileNetV1() elif args.arch.startswith('mobilenetv2'): model = MobileNetV2() else: model = models.__dict__[args.arch]() if args.pretrainedmodeler: if os.path.isfile(args.pretrainedmodeler): print("=> loading model '{}'".format(args.pretrainedmodeler)) new_state_dict = OrderedDict() state_dict = torch.load(args.pretrainedmodeler) for k, v in state_dict.items(): name = k[7:] new_state_dict[name] = v model.load_state_dict(new_state_dict) print(model) if not args.distributed: if args.arch.startswith('alexnet') or args.arch.startswith('vgg'): model.features = torch.nn.DataParallel(model.features) model.cuda() else: model = torch.nn.DataParallel(model).cuda() else: model.cuda() model = torch.nn.parallel.DistributedDataParallel(model) torch.save(model, f"saves/model_after_retraining.ptmodel") if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) if args.evaluate: if os.path.isfile(args.evaluate): print("=> loading model '{}'".format(args.evaluate)) model.load_state_dict(torch.load(args.evaluate)) else: print("=> no model found at '{}'".format(args.evaluate)) # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})" .format(args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True # Data loading code traindir = os.path.join(args.data, 'train') valdir = os.path.join(args.data, 'val') normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_dataset = datasets.ImageFolder( traindir, transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) else: train_sampler = None train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), num_workers=args.workers, pin_memory=True, sampler=train_sampler) val_loader = torch.utils.data.DataLoader( datasets.ImageFolder(valdir, transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ])), batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) if args.evaluate: validate(val_loader, model, criterion) return for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) adjust_learning_rate(optimizer, epoch) # train for one epoch train(train_loader, model, criterion, optimizer, epoch) # evaluate on validation set prec1 = validate(val_loader, model, criterion) # remember best prec@1 and save checkpoint is_best = prec1 > best_prec1 best_prec1 = max(prec1, best_prec1) save_checkpoint({ 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer' : optimizer.state_dict(), }, is_best) print("--- Before pruning ---") util.print_nonzeros(model)
def inference(model): util.print_nonzeros(model, f"{args.save_dir}/{args.log}") accuracy = util.validate(val_loader, model, args) print(accuracy)
def main(): global best_acc start_epoch = args.start_epoch # start from epoch 0 or last checkpoint epoch # Data print('==> Preparing dataset %s' % args.dataset) transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) if args.dataset == 'cifar10': dataloader = datasets.CIFAR10 num_classes = 10 else: dataloader = datasets.CIFAR100 num_classes = 100 trainset = dataloader(root='./data', train=True, download=True, transform=transform_train) trainloader = data.DataLoader(trainset, batch_size=args.train_batch, shuffle=True, num_workers=args.workers) testset = dataloader(root='./data', train=False, download=False, transform=transform_test) testloader = data.DataLoader(testset, batch_size=args.test_batch, shuffle=False, num_workers=args.workers) # Model print("==> creating model '{}'".format(args.arch)) if args.arch.startswith('resnext'): model = models.__dict__[args.arch]( cardinality=args.cardinality, num_classes=num_classes, depth=args.depth, widen_factor=args.widen_factor, dropRate=args.drop, ) elif args.arch.startswith('densenet'): model = models.__dict__[args.arch]( num_classes=num_classes, depth=args.depth, growthRate=args.growthRate, compressionRate=args.compressionRate, dropRate=args.drop, ) elif args.arch.startswith('wrn'): model = models.__dict__[args.arch]( num_classes=num_classes, depth=args.depth, widen_factor=args.widen_factor, dropRate=args.drop, ) elif args.arch.endswith('resnet'): model = models.__dict__[args.arch]( num_classes=num_classes, depth=args.depth, block_name=args.block_name, ) else: model = models.__dict__[args.arch](num_classes=num_classes) model = torch.nn.DataParallel(model).cuda() cudnn.benchmark = True print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0)) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # Resume title = 'cifar-10-' + args.arch model.load_state_dict( torch.load(os.path.join(args.resume, args.model + '.pth'))) logger = Logger(os.path.join(save_path, 'log.txt'), title=title) logger.set_names([ 'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.', 'Valid Acc.' ]) print('model loaded') device = torch.device("cuda") # Initial training print("--- Pruning ---") masks = {} for name, p in model.named_parameters(): if 'weight' in name: tensor = p.data.cpu().numpy() threshold = args.sensitivity new_mask = np.where(abs(tensor) < threshold, 0., tensor) mask = np.where(abs(tensor) < threshold, 0., 1.) masks[name] = torch.from_numpy(mask).float().to(device) p.data = torch.from_numpy(new_mask).to(device) util.print_nonzeros(model) print('Pruned model evaluation...') test_loss, test_acc = test(testloader, model, criterion, start_epoch, use_cuda) print(' Test Loss: %.8f, Test Acc: %.2f' % (test_loss, test_acc)) if args.evaluate: return best_prec1 = test_acc torch.save(model.state_dict(), os.path.join(save_path, 'finetuned.pth')) print("--- Finetuning ---") # Train and val for epoch in range(start_epoch, args.epochs): adjust_learning_rate(optimizer, epoch) print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, args.epochs, state['lr'])) train_loss, train_acc = train(trainloader, model, criterion, optimizer, epoch, use_cuda, masks) test_loss, test_acc = test(testloader, model, criterion, epoch, use_cuda) # append logger file logger.append( [state['lr'], train_loss, test_loss, train_acc, test_acc]) # save model is_best = test_acc > best_acc best_acc = max(test_acc, best_acc) if is_best: torch.save(model.state_dict(), os.path.join(save_path, 'finetuned.pth')) logger.close() #logger.plot() #savefig(os.path.join(save_path, 'log.eps')) print("--- Evaluating ---") model.load_state_dict(torch.load(os.path.join(save_path, 'finetuned.pth'))) test_loss, test_acc = test(testloader, model, criterion, start_epoch, use_cuda) print(' Test Loss: %.8f, Test Acc: %.2f' % (test_loss, test_acc)) print('Best acc:') print(best_acc)
def main(): if args.seed is not None: random.seed(args.seed) torch.manual_seed(args.seed) cudnn.deterministic = True warnings.warn('You have chosen to seed training. ' 'This will turn on the CUDNN deterministic setting, ' 'which can slow down your training considerably! ' 'You may see unexpected behavior when restarting ' 'from checkpoints.') if args.gpu is not None: warnings.warn('You have chosen a specific GPU. This will completely ' 'disable data parallelism.') args.distributed = args.world_size > 1 if args.distributed: dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size) # create model if args.pretrained: print("=> using pre-trained model '{}'".format(args.arch)) model = models.__dict__[args.arch](pretrained=True) else: print("=> creating model '{}'".format(args.arch)) model = models.__dict__[args.arch]() if args.gpu is not None: model = model.cuda(args.gpu) elif args.distributed: model.cuda() model = torch.nn.parallel.DistributedDataParallel(model) else: if args.arch.startswith('alexnet') or args.arch.startswith('vgg'): if args.arch.startswith('alexnetv2') and args.pretrained: pass else: model.features = torch.nn.DataParallel(model.features) model.cuda() else: model = torch.nn.DataParallel(model).cuda() input_size = 224 # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda(args.gpu) optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) if args.pretrained is not True: model.load_state_dict(torch.load(os.path.join(args.resume, args.model+'.pth'))) print('model loaded') cudnn.benchmark = True # Data loading code traindir = os.path.join(args.data, 'train') valdir = os.path.join(args.data, 'val') normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_dataset = datasets.ImageFolder( traindir, transforms.Compose([ transforms.RandomResizedCrop(input_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) else: train_sampler = None train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), num_workers=args.workers, pin_memory=True, sampler=train_sampler) val_loader = torch.utils.data.DataLoader( datasets.ImageFolder(valdir, transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(input_size), transforms.ToTensor(), normalize, ])), batch_size=256, shuffle=False, num_workers=args.workers, pin_memory=True) device = torch.device("cuda") # Initial pruning print("--- Pruning ---") masks = {} for name, p in model.named_parameters(): if 'weight' in name: tensor = p.data.cpu().numpy() threshold = args.sensitivity new_mask = np.where(abs(tensor) < threshold, 0, tensor) mask = np.where(abs(tensor) < threshold, 0., 1.) masks[name] = torch.from_numpy(mask).float().to(device) p.data = torch.from_numpy(new_mask).to(device) util.print_nonzeros(model) print('Pruned model evaluation...') prec1, prec5 = validate(val_loader, model, criterion) if args.evaluate: return best_prec1 = prec5 torch.save(model.state_dict(), os.path.join(save_path, 'finetuned.pth')) print("--- Finetuning ---") curves = np.zeros(((args.epochs-args.start_epoch)*(len(train_loader)//args.print_freq),2)) valid = np.zeros(((args.epochs-args.start_epoch),3)) step = 0 for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) adjust_learning_rate(optimizer, epoch) # train for one epoch curves,step = train(train_loader, model, criterion, optimizer, epoch, curves, step, masks) # evaluate on validation set valid[epoch, 0] = epoch valid[epoch, 1], valid[epoch, 2] = validate(val_loader, model, criterion) prec5 = valid[epoch, 2] # plot training curve np.savetxt(os.path.join(save_path, 'curves.dat'), curves) clr1 = (0.5, 0., 0.) clr2 = (0.0, 0.5, 0.) fig, ax1 = plt.subplots() ax1.set_xlabel('steps') ax1.set_ylabel('Loss', color=clr1) ax1.tick_params(axis='y', colors=clr1) start = 0 end = step markersize = 12 coef = 2. ax1.plot(curves[start:end, 0], curves[start:end, 1], '--', color=[c*coef for c in clr1], markersize=markersize) #ax2.set_ylim(bottom=20, top=100) ax1.legend(('Train loss'), loc='lower right') fig.savefig(os.path.join(save_path, 'loss-vs-steps.pdf')) # plot validation curve np.savetxt(os.path.join(save_path, 'valid.dat'), valid) fig3, ax5 = plt.subplots() ax6 = ax5.twinx() ax5.set_xlabel('epochs') ax5.set_ylabel('Acc@1', color=clr1) ax5.tick_params(axis='y', colors=clr1) ax6.set_ylabel('Acc@5', color=clr2) ax6.tick_params(axis='y', colors=clr2) start = 0 end = epoch+1 markersize = 12 coef = 2. ax5.plot(valid[start:end, 0], valid[start:end, 1], '--', color=[c*coef for c in clr1], markersize=markersize) ax6.plot(valid[start:end, 0], valid[start:end, 2], '-', color=[c*coef for c in clr2], markersize=markersize) #ax2.set_ylim(bottom=20, top=100) ax5.legend(('Acc@1'), loc='lower right') ax6.legend(('Acc@5'), loc='lower left') fig3.savefig(os.path.join(save_path, 'accuracy-vs-epochs.pdf')) if prec5 > best_prec1: torch.save(model.state_dict(), os.path.join(save_path, 'finetuned.pth')) best_prec1 = prec5 print('New best performance') print("--- Evaluating ---") model.load_state_dict(torch.load(os.path.join(save_path, 'finetuned.pth'))) prec1 = validate(val_loader, model, criterion)
test_loss /= len(test_loader.dataset) accuracy = 100. * correct / len(test_loader.dataset) print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)') return accuracy # Initial training print("--- Initial training ---") util.log(args.log, "--- Initial training ---") train(100) accuracy = test() util.log(args.log, f"initial_accuracy {accuracy}") torch.save(model, f"saves/initial_model.ptmodel") print("--- Before pruning ---") util.log(args.log, "--- Before pruning ---") util.log(args.log, util.print_nonzeros(model)) original_accuracy = accuracy loop_cnt = 0 minus_cnt = 0 start = time.time() while (time.time() - start) / 3600 < 2: print("Iteration", loop_cnt) util.log(args.log, f"Iteration {loop_cnt}") loop_cnt += 1 # Pruning model.prune_by_std(args.sensitivity) accuracy = test() util.log(args.log, f"accuracy_after_pruning {accuracy}") print("--- After pruning ---")