def packing_job_ext_info(job_lsit_DO): """ Packing additional information of the job into the job_list_DO(JobListDO) """ ext_info = sqllite_agent.execute(ScrapydJobExtInfoSQLSet.SELECT_BY_ID, (job_lsit_DO.job_id, )) if ext_info is None or len(ext_info) <= 0: return ext_info = ext_info[0] job_lsit_DO.args = ext_info[1] job_lsit_DO.priority = ext_info[2] job_lsit_DO.creation_time = ext_info[3] job_lsit_DO.logs_name = str_to_list(ext_info[4], ',') job_lsit_DO.logs_url = str_to_list(ext_info[5], ',')
def main(config): # init logger classes = { 'cifar10': 10, 'cifar100': 100, 'mnist': 10, 'tiny_imagenet': 200, 'imagenet': 1000 } logger, writer = init_logger(config) # build model model = models.__dict__[config.network]() mb = ModelBase(config.network, config.depth, config.dataset, model) mb.cuda() # preprocessing # ====================================== fetch configs ====================================== ckpt_path = config.checkpoint_dir num_iterations = config.iterations target_ratio = config.target_ratio normalize = config.normalize # ====================================== fetch exception ====================================== exception = get_exception_layers(mb.model, str_to_list(config.exception, ',', int)) logger.info('Exception: ') for idx, m in enumerate(exception): logger.info(' (%d) %s' % (idx, m)) # ====================================== fetch training schemes ====================================== ratio = 1 - (1 - target_ratio)**(1.0 / num_iterations) learning_rates = str_to_list(config.learning_rate, ',', float) weight_decays = str_to_list(config.weight_decay, ',', float) training_epochs = str_to_list(config.epoch, ',', int) logger.info( 'Normalize: %s, Total iteration: %d, Target ratio: %.2f, Iter ratio %.4f.' % (normalize, num_iterations, target_ratio, ratio)) logger.info('Basic Settings: ') for idx in range(len(learning_rates)): logger.info(' %d: LR: %.5f, WD: %.5f, Epochs: %d' % (idx, learning_rates[idx], weight_decays[idx], training_epochs[idx])) # ====================================== get dataloader ====================================== normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_dataset = datasets.ImageFolder( config.traindir, transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=250, shuffle=True, num_workers=16, pin_memory=True, sampler=None) # ====================================== start pruning ====================================== for iteration in range(num_iterations): logger.info( '** Target ratio: %.4f, iter ratio: %.4f, iteration: %d/%d.' % (target_ratio, ratio, iteration, num_iterations)) assert num_iterations == 1 print("=> Applying weight initialization.") mb.model.apply(weights_init) masks = GraSP(mb.model, ratio, trainloader, 'cuda', num_classes=classes[config.dataset], samples_per_class=config.samples_per_class, num_iters=config.get('num_iters', 1)) # ========== register mask ================== mb.masks = masks # ========== save pruned network ============ logger.info('Saving..') state = { 'net': mb.model, 'acc': -1, 'epoch': -1, 'args': config, 'mask': mb.masks, 'ratio': mb.get_ratio_at_each_layer() } path = os.path.join( ckpt_path, 'prune_%s_%s%s_r%s_it%d.pth.tar' % (config.dataset, config.network, config.depth, config.target_ratio, iteration)) torch.save(state, path) # ========== print pruning details ============ logger.info('**[%d] Mask and training setting: ' % iteration) print_mask_information(mb, logger)
def main(config, args): # init logger classes = { 'cifar10': 10, 'cifar100': 100, 'mnist': 10, 'tiny_imagenet': 200 } logger, writer = init_logger(config, args) best_acc_vec = [] test_acc_vec_vec = [] for n_runs in range(1): if args.sigma_w2 != None and n_runs != 0: break # build model model = get_network(config.network, config.depth, config.dataset, use_bn=config.get('use_bn', args.bn), scaled=args.scaled_init, act=args.act) mask = None mb = ModelBase(config.network, config.depth, config.dataset, model) mb.cuda() if mask is not None: mb.register_mask(mask) ratio_vec_ = print_mask_information(mb, logger) # preprocessing # ====================================== get dataloader ====================================== trainloader, testloader = get_dataloader(config.dataset, config.batch_size, 256, 4) # ====================================== fetch configs ====================================== ckpt_path = config.checkpoint_dir num_iterations = config.iterations if args.target_ratio == None: target_ratio = config.target_ratio else: target_ratio = args.target_ratio normalize = config.normalize # ====================================== fetch exception ====================================== exception = get_exception_layers( mb.model, str_to_list(config.exception, ',', int)) logger.info('Exception: ') for idx, m in enumerate(exception): logger.info(' (%d) %s' % (idx, m)) # ====================================== fetch training schemes ====================================== ratio = 1 - (1 - target_ratio)**(1.0 / num_iterations) learning_rates = str_to_list(config.learning_rate, ',', float) weight_decays = str_to_list(config.weight_decay, ',', float) training_epochs = str_to_list(config.epoch, ',', int) logger.info( 'Normalize: %s, Total iteration: %d, Target ratio: %.2f, Iter ratio %.4f.' % (normalize, num_iterations, target_ratio, ratio)) logger.info('Basic Settings: ') for idx in range(len(learning_rates)): logger.info(' %d: LR: %.5f, WD: %.5f, Epochs: %d' % (idx, learning_rates[idx], weight_decays[idx], training_epochs[idx])) # ====================================== start pruning ====================================== iteration = 0 for _ in range(1): logger.info( '** Target ratio: %.4f, iter ratio: %.4f, iteration: %d/%d.' % (target_ratio, ratio, 1, num_iterations)) # mb.model.apply(weights_init) print('#' * 40) print('USING {} INIT SCHEME'.format(args.init)) print('#' * 40) if args.init == 'kaiming_xavier': mb.model.apply(weights_init_kaiming_xavier) elif args.init == 'kaiming': if args.act == 'relu' or args.act == 'elu': mb.model.apply(weights_init_kaiming_relu) elif args.act == 'tanh': mb.model.apply(weights_init_kaiming_tanh) elif args.init == 'xavier': mb.model.apply(weights_init_xavier) elif args.init == 'EOC': mb.model.apply(weights_init_EOC) elif args.init == 'ordered': def weights_init_ord(m): if isinstance(m, nn.Conv2d): ord_weights(m.weight, sigma_w2=args.sigma_w2) if m.bias is not None: ord_bias(m.bias) elif isinstance(m, nn.Linear): ord_weights(m.weight, sigma_w2=args.sigma_w2) if m.bias is not None: ord_bias(m.bias) elif isinstance(m, nn.BatchNorm2d): # Note that BN's running_var/mean are # already initialized to 1 and 0 respectively. if m.weight is not None: m.weight.data.fill_(1.0) if m.bias is not None: m.bias.data.zero_() mb.model.apply(weights_init_ord) else: raise NotImplementedError print("=> Applying weight initialization(%s)." % config.get('init_method', 'kaiming')) print("Iteration of: %d/%d" % (iteration, num_iterations)) if config.pruner == 'SNIP': print('=> Using SNIP') masks, scaled_masks = SNIP( mb.model, ratio, trainloader, 'cuda', num_classes=classes[config.dataset], samples_per_class=config.samples_per_class, num_iters=config.get('num_iters', 1), scaled_init=args.scaled_init) elif config.pruner == 'GraSP': print('=> Using GraSP') masks, scaled_masks = GraSP( mb.model, ratio, trainloader, 'cuda', num_classes=classes[config.dataset], samples_per_class=config.samples_per_class, num_iters=config.get('num_iters', 1), scaled_init=args.scaled_init) iteration = 0 ################################################################################ _masks = None _masks_scaled = None if not args.bn: # build model that has the same weights as the pruned network but with BN now ! model2 = get_network(config.network, config.depth, config.dataset, use_bn=config.get('use_bn', True), scaled=args.scaled_init, act=args.act) weights_temp = [] for layer_old in mb.model.modules(): if isinstance(layer_old, nn.Conv2d) or isinstance( layer_old, nn.Linear): weights_temp.append(layer_old.weight) idx = 0 for layer_new in model2.modules(): if isinstance(layer_new, nn.Conv2d) or isinstance( layer_new, nn.Linear): layer_new.weight.data = weights_temp[idx] idx += 1 # Creating a base model with BN included now mb = ModelBase(config.network, config.depth, config.dataset, model2) mb.cuda() _masks = dict() _masks_scaled = dict() layer_keys_new = [] for layer in (mb.model.modules()): if isinstance(layer, nn.Conv2d) or isinstance( layer, nn.Linear): layer_keys_new.append(layer) for new_keys, old_keys in zip(layer_keys_new, masks.keys()): _masks[new_keys] = masks[old_keys] if args.scaled_init: _masks_scaled[new_keys] = scaled_masks[old_keys] ################################################################################ if _masks == None: _masks = masks _masks_scaled = scaled_masks # ========== register mask ================== mb.register_mask(_masks) ## ========== debugging ================== if args.scaled_init: if config.network == 'vgg': print('scaling VGG') mb.scaling_weights(_masks_scaled) # ========== save pruned network ============ logger.info('Saving..') state = { 'net': mb.model, 'acc': -1, 'epoch': -1, 'args': config, 'mask': mb.masks, 'ratio': mb.get_ratio_at_each_layer() } path = os.path.join( ckpt_path, 'prune_%s_%s%s_r%s_it%d.pth.tar' % (config.dataset, config.network, config.depth, target_ratio, iteration)) torch.save(state, path) # ========== print pruning details ============ logger.info('**[%d] Mask and training setting: ' % iteration) ratio_vec_ = print_mask_information(mb, logger) logger.info(' LR: %.5f, WD: %.5f, Epochs: %d' % (learning_rates[iteration], weight_decays[iteration], training_epochs[iteration])) results_path = config.summary_dir + args.init + '_sp' + str( args.target_ratio).replace('.', '_') if args.scaled_init: results_path += '_scaled' if args.bn: results_path += '_bn' if args.sigma_w2 != None and args.init == 'ordered': results_path += '_sgw2{}'.format(args.sigma_w2).replace( '.', '_') results_path += '_' + args.act + '_' + str(config.depth) print('saving the ratios') print(results_path) if not os.path.isdir(results_path): os.mkdir(results_path) np.save(results_path + '/ratios_pruned{}'.format(args.seed_tiny), np.array(ratio_vec_)) # if args.sigma_w2 != None: # break # ========== finetuning ======================= best_acc, test_acc_vec = train_once( mb=mb, net=mb.model, trainloader=trainloader, testloader=testloader, writer=writer, config=config, ckpt_path=ckpt_path, learning_rate=learning_rates[iteration], weight_decay=weight_decays[iteration], num_epochs=training_epochs[iteration], iteration=iteration, logger=logger, args=args) best_acc_vec.append(best_acc) test_acc_vec_vec.append(test_acc_vec) np.save(results_path + '/best_acc{}'.format(args.seed_tiny), np.array(best_acc_vec)) np.save(results_path + '/test_acc{}'.format(args.seed_tiny), np.array(test_acc_vec_vec))
def main(config): stats = {} device = 'cuda' criterion = torch.nn.CrossEntropyLoss() # config = init_config() if config is None else config logger, writer = init_summary_writer(config) trainloader, testloader = init_dataloader(config) net, bottleneck_net = init_network(config, logger, device) pruner = init_pruner(net, bottleneck_net, config, writer, logger) # start pruning epochs = str_to_list(config.epoch, ',', int) learning_rates = str_to_list(config.learning_rate, ',', float) weight_decays = str_to_list(config.weight_decay, ',', float) ratios = str_to_list(config.ratio, ',', float) fisher_type = config.fisher_type # empirical|true fisher_mode = config.fisher_mode # eigen|full|diagonal normalize = config.normalize prune_mode = config.prune_mode # one-pass | iterative fix_rotation = config.get('fix_rotation', True) assert (len(epochs) == len(learning_rates) and len(learning_rates) == len(weight_decays) and len(weight_decays) == len(ratios)) total_parameters = count_parameters(net.train()) for it in range(len(epochs)): epoch = epochs[it] lr = learning_rates[it] wd = weight_decays[it] ratio = ratios[it] logger.info('-' * 120) logger.info('** [%d], Ratio: %.2f, epoch: %d, lr: %.4f, wd: %.4f' % (it, ratio, epoch, lr, wd)) logger.info( 'Reinit: %s, Fisher_mode: %s, fisher_type: %s, normalize: %s, fix_rotation: %s.' % (config.re_init, fisher_mode, fisher_type, normalize, fix_rotation)) pruner.fix_rotation = fix_rotation # conduct pruning cfg = pruner.make_pruned_model(trainloader, criterion=criterion, device=device, fisher_type=fisher_type, prune_ratio=ratio, normalize=normalize, re_init=config.re_init) # for tracking the best accuracy compression_ratio, unfair_ratio, all_numel, rotation_numel = compute_ratio( pruner.model, total_parameters, fix_rotation, logger) if config.dataset == 'tiny_imagenet': total_flops, rotation_flops = print_model_param_flops(pruner.model, 64, cuda=True) else: total_flops, rotation_flops = print_model_param_flops(pruner.model, 32, cuda=True) train_loss_pruned, train_acc_pruned = pruner.test_model( trainloader, criterion, device) test_loss_pruned, test_acc_pruned = pruner.test_model( testloader, criterion, device) # write results logger.info('Before: Accuracy: %.2f%%(train), %.2f%%(test).' % (train_acc_pruned, test_acc_pruned)) logger.info(' Loss: %.2f (train), %.2f (test).' % (train_loss_pruned, test_loss_pruned)) test_loss_finetuned, test_acc_finetuned = pruner.fine_tune_model( trainloader=trainloader, testloader=testloader, criterion=criterion, optim=optim, learning_rate=lr, weight_decay=wd, nepochs=epoch) train_loss_finetuned, train_acc_finetuned = pruner.test_model( trainloader, criterion, device) logger.info('After: Accuracy: %.2f%%(train), %.2f%%(test).' % (train_acc_finetuned, test_acc_finetuned)) logger.info(' Loss: %.2f (train), %.2f (test).' % (train_loss_finetuned, test_loss_finetuned)) # save model stat = { 'total_flops': total_flops, 'rotation_flops': rotation_flops, 'it': it, 'prune_ratio': ratio, 'cr': compression_ratio, 'unfair_cr': unfair_ratio, 'all_params': all_numel, 'rotation_params': rotation_numel, 'prune/train_loss': train_loss_pruned, 'prune/train_acc': train_acc_pruned, 'prune/test_loss': test_loss_pruned, 'prune/test_acc': test_acc_pruned, 'finetune/train_loss': train_loss_finetuned, 'finetune/test_loss': test_loss_finetuned, 'finetune/train_acc': train_acc_finetuned, 'finetune/test_acc': test_acc_finetuned } save_model(config, it, pruner, cfg, stat) stats[it] = stat if prune_mode == 'one_pass': del net del pruner net, bottleneck_net = init_network(config, logger, device) pruner = init_pruner(net, bottleneck_net, config, writer, logger) pruner.iter = it with open(os.path.join(config.summary_dir, 'stats.json'), 'w') as f: json.dump(stats, f)
def main(config): # init logger classes = { 'cifar10': 10, 'cifar100': 100, 'mnist': 10, 'tiny_imagenet': 200 } logger, writer = init_logger(config) # build model model = get_network(config.network, config.depth, config.dataset, use_bn=config.get('use_bn', True)) mask = None mb = ModelBase(config.network, config.depth, config.dataset, model) mb.cuda() if mask is not None: mb.register_mask(mask) print_mask_information(mb, logger) # preprocessing # ====================================== get dataloader ====================================== trainloader, testloader = get_dataloader(config.dataset, config.batch_size, 256, 4, root='/home/wzn/PycharmProjects/GraSP/data') # ====================================== fetch configs ====================================== ckpt_path = config.checkpoint_dir num_iterations = config.iterations target_ratio = config.target_ratio normalize = config.normalize # ====================================== fetch exception ====================================== # exception = get_exception_layers(mb.model, str_to_list(config.exception, ',', int)) # logger.info('Exception: ') # # for idx, m in enumerate(exception): # logger.info(' (%d) %s' % (idx, m)) # ====================================== fetch training schemes ====================================== ratio = 1 - (1 - target_ratio) ** (1.0 / num_iterations) learning_rates = str_to_list(config.learning_rate, ',', float) weight_decays = str_to_list(config.weight_decay, ',', float) training_epochs = str_to_list(config.epoch, ',', int) logger.info('Normalize: %s, Total iteration: %d, Target ratio: %.2f, Iter ratio %.4f.' % (normalize, num_iterations, target_ratio, ratio)) logger.info('Basic Settings: ') for idx in range(len(learning_rates)): logger.info(' %d: LR: %.5f, WD: %.5f, Epochs: %d' % (idx, learning_rates[idx], weight_decays[idx], training_epochs[idx])) # ====================================== start pruning ====================================== iteration = 0 for _ in range(1): # logger.info('** Target ratio: %.4f, iter ratio: %.4f, iteration: %d/%d.' % (target_ratio, # ratio, # 1, # num_iterations)) mb.model.apply(weights_init) print("=> Applying weight initialization(%s)." % config.get('init_method', 'kaiming')) # print("Iteration of: %d/%d" % (iteration, num_iterations)) # masks = GraSP(mb.model, ratio, trainloader, 'cuda', # num_classes=classes[config.dataset], # samples_per_class=config.samples_per_class, # num_iters=config.get('num_iters', 1)) # iteration = 0 # print('=> Using GraSP') # # ========== register mask ================== # mb.register_mask(masks) # # ========== save pruned network ============ # logger.info('Saving..') # state = { # 'net': mb.model, # 'acc': -1, # 'epoch': -1, # 'args': config, # 'mask': mb.masks, # 'ratio': mb.get_ratio_at_each_layer() # } # path = os.path.join(ckpt_path, 'prune_%s_%s%s_r%s_it%d.pth.tar' % (config.dataset, # config.network, # config.depth, # config.target_ratio, # iteration)) # torch.save(state, path) # # ========== print pruning details ============ # logger.info('**[%d] Mask and training setting: ' % iteration) # print_mask_information(mb, logger) # logger.info(' LR: %.5f, WD: %.5f, Epochs: %d' % # (learning_rates[iteration], weight_decays[iteration], training_epochs[iteration])) # ========== finetuning ======================= train_once(mb=mb, net=mb.model, trainloader=trainloader, testloader=testloader, writer=writer, config=config, ckpt_path=ckpt_path, learning_rate=learning_rates[iteration], weight_decay=weight_decays[iteration], num_epochs=training_epochs[iteration], iteration=iteration, logger=logger)
def main(config): stats = {} if config.data_distributed: torch.distributed.init_process_group(backend="nccl") device = torch.device('cuda:0,1') criterion = torch.nn.CrossEntropyLoss() logger, writer = init_summary_writer(config) trainloader, testloader = init_dataloader(config) if config.data_distributed: trainset, testset = trainloader, testloader train_sampler = torch.utils.data.distributed.DistributedSampler( trainset) test_sampler = torch.utils.data.distributed.DistributedSampler(testset) trainloader = torch.utils.data.DataLoader( trainset, config.batch_size, False, num_workers=config.num_workers, pin_memory=True, drop_last=True, sampler=train_sampler) testloader = torch.utils.data.DataLoader( testset, config.batch_size, False, num_workers=config.num_workers, pin_memory=True, drop_last=True, sampler=test_sampler) hess_data = [] if config.dataset == 'imagenet': hessianloader = get_hessianloader(config.dataset, 64) num_batch = config.hessian_batch_size // 64 i = 0 for data, label in hessianloader: i += 1 hess_data.append((data, label)) if i == num_batch: break else: hessianloader = get_hessianloader(config.dataset, config.hessian_batch_size) for data, label in hessianloader: hess_data = (data, label) net, bottleneck_net = init_network(config, logger, device, config.dataset == "imagenet") pruner = init_pruner(net, bottleneck_net, config, writer, logger) # total FLOPs calculation if config.dataset == 'tiny_imagenet': total_flops, _ = compute_model_param_flops(pruner.model, 64, cuda=True) elif config.dataset == 'imagenet': total_flops, _ = compute_model_param_flops(pruner.model, 224, cuda=True) else: total_flops, _ = compute_model_param_flops(pruner.model, 32, cuda=True) # start pruning epochs = str_to_list(config.epoch, ',', int) learning_rates = str_to_list(config.learning_rate, ',', float) weight_decays = str_to_list(config.weight_decay, ',', float) ratios = str_to_list(config.ratio, ',', float) fisher_type = config.fisher_type # empirical|true fisher_mode = config.fisher_mode # eigen|full|diagonal normalize = config.normalize prune_mode = config.prune_mode # one-pass | iterative fix_rotation = config.get('fix_rotation', True) assert (len(epochs) == len(learning_rates) and len(learning_rates) == len(weight_decays) and len(weight_decays) == len(ratios)) total_parameters = count_parameters(net) for it in range(len(epochs)): epoch = epochs[it] lr = learning_rates[it] wd = weight_decays[it] ratio = ratios[it] logger.info('-' * 120) logger.info('** [%d], Ratio: %.2f, epoch: %d, lr: %.4f, wd: %.4f' % (it, ratio, epoch, lr, wd)) logger.info( 'Reinit: %s, Fisher_mode: %s, fisher_type: %s, normalize: %s, fix_rotation: %s.' % (config.re_init, fisher_mode, fisher_type, normalize, fix_rotation)) pruner.fix_rotation = fix_rotation # test pretrained model if config.init_test: train_loss_pruned, train_acc_pruned, top5_acc = pruner.test_model( trainloader, criterion, device) test_loss_pruned, test_acc_pruned, top5_acc = pruner.test_model( testloader, criterion, device) logger.info('Pretrain: Accuracy: %.2f%%(train), %.2f%%(test).' % (train_acc_pruned, test_acc_pruned)) logger.info(' Loss: %.2f (train), %.2f (test).' % (train_loss_pruned, test_loss_pruned)) # conduct pruning if 'hessian' not in config.fisher_mode: cfg = pruner.make_pruned_model(trainloader, criterion=criterion, device=device, fisher_type=fisher_type, prune_ratio=ratio, normalize=normalize, re_init=config.re_init) else: cfg = pruner.make_pruned_model(hess_data, criterion=criterion, device=device, fisher_type=fisher_type, prune_ratio=ratio, normalize=normalize, re_init=config.re_init, n_v=config.nv) print(pruner.model) # for tracking the best accuracy compression_ratio, unfair_ratio, all_numel, rotation_numel = compute_ratio( pruner.model, total_parameters, fix_rotation, logger) if config.dataset == 'tiny_imagenet': remained_flops, rotation_flops = compute_model_param_flops( pruner.model, 64, cuda=True) logger.info( ' + Remained FLOPs: %.4fG(%.2f%%), Total FLOPs: %.4fG' % (remained_flops / 1e9, 100. * remained_flops / total_flops, total_flops / 1e9)) elif config.dataset == 'imagenet': remained_flops, rotation_flops = compute_model_param_flops( pruner.model, 224, cuda=True) logger.info( ' + Remained FLOPs: %.4fG(%.2f%%), Total FLOPs: %.4fG' % (remained_flops / 1e9, 100. * remained_flops / total_flops, total_flops / 1e9)) else: remained_flops, rotation_flops = compute_model_param_flops( pruner.model, 32, cuda=True) logger.info( ' + Remained FLOPs: %.4fG(%.2f%%), Total FLOPs: %.4fG' % (remained_flops / 1e9, 100. * remained_flops / total_flops, total_flops / 1e9)) logger.info(f"Total Flops: {remained_flops}") test_loss_pruned, test_acc_pruned, top5_acc = pruner.test_model( testloader, criterion, device) if config.dataset != 'imagenet': time_loader = get_hessianloader(config.dataset, 1) run_time = pruner.speed_model(time_loader, criterion, device) logger.info(f"Total Run Time: {run_time}") test_loss_finetuned, test_acc_finetuned = pruner.fine_tune_model( trainloader=trainloader, testloader=testloader, criterion=criterion, optim=optim, learning_rate=lr, weight_decay=wd, nepochs=epoch) train_loss_finetuned, train_acc_finetuned, top5_acc = pruner.test_model( trainloader, criterion, device) logger.info( f'After {config.dataset, config.network, config.depth}: Accuracy: %.2f%%(train), %.2f%%.' % (train_acc_finetuned, test_acc_finetuned)) logger.info(' Loss: %.2f (train), %.2f .' % (train_loss_finetuned, test_loss_finetuned)) stat = { 'total_flops': total_flops, 'rotation_flops': rotation_flops, 'flops_remained': float(100. * remained_flops / total_flops), 'it': it, 'prune_ratio': ratio, 'cr': compression_ratio, 'unfair_cr': unfair_ratio, 'all_params': all_numel, 'rotation_params': rotation_numel, 'prune/test_loss': test_loss_pruned, 'prune/test_acc': test_acc_pruned, 'finetune/train_loss': train_loss_finetuned, 'finetune/test_loss': test_loss_finetuned, 'finetune/train_acc': train_acc_finetuned, 'finetune/test_acc': test_acc_finetuned } print('saving checkpoint') save_model(config, it, pruner, cfg, stat) stats[it] = stat if prune_mode == 'one_pass': print('one_pass') del net del pruner net, bottleneck_net = init_network(config, logger, device, config.dataset == "imagenet") pruner = init_pruner(net, bottleneck_net, config, writer, logger) pruner.iter = it with open( os.path.join(config.saving_log, f'stats_{running_time}.json'), 'w') as f: json.dump(stats, f) if prune_mode != 'one_pass': with open(os.path.join(config.saving_log, f'stats{it}.json'), 'w') as f: json.dump(stats, f)