run_manager.save_config(print_info=True)
# load checkpoints
if args.base_path!=None:
weight_path = args.base_path+'/checkpoint/model_best.pth.tar'
if args.resume:
run_manager.load_model()
if args.train and run_manager.best_acc == 0:
loss, acc1, acc5 = run_manager.validate(is_test=True, return_top5=True)
run_manager.best_acc = acc1
elif weight_path!=None and os.path.isfile(weight_path):
assert net_origin != None, "original network is None"
net_origin.load_state_dict(torch.load(weight_path)['state_dict'])
net_origin = net_origin.module
if args.model == 'resnet18':
run_manager.reset_model(ResNet_ImageNet(num_classes=run_config.data_provider.n_classes, cfg=eval(args.cfg), depth=18), net_origin.cpu())
elif args.model == 'resnet34':
run_manager.reset_model(ResNet_ImageNet(num_classes=run_config.data_provider.n_classes, cfg=eval(args.cfg), depth=34), net_origin.cpu())
elif args.model == 'resnet50':
run_manager.reset_model(ResNet_ImageNet(num_classes=run_config.data_provider.n_classes, cfg=eval(args.cfg), depth=50), net_origin.cpu())
elif args.model == 'mobilenet':
run_manager.reset_model(MobileNet(num_classes=run_config.data_provider.n_classes, cfg=eval(args.cfg)), net_origin.cpu())
elif args.model == 'mobilenetv2':
run_manager.reset_model(MobileNetV2(num_classes=run_config.data_provider.n_classes, cfg=eval(args.cfg)), net_origin.cpu())
elif args.model == 'vgg':
run_manager.reset_model(VGG_CIFAR(cfg=eval(args.cfg), cutout=True), net_origin.cpu())
elif args.model == 'resnet56':
run_manager.reset_model(ResNet_CIFAR(cfg=eval(args.cfg), depth=56, num_classes=run_config.data_provider.n_classes, cutout=True), net_origin.cpu())
elif args.model== 'resnet110':
run_manager.reset_model(ResNet_CIFAR(cfg=eval(args.cfg), depth=110, num_classes=run_config.data_provider.n_classes, cutout=True), net_origin.cpu())
else:
'prune')
current_best = np.array(fitness).max()
if current_best > best_acc:
best_acc = current_best
best_cfg = cfgs[np.array(fitness).argsort()[-1]]
best_cfg_raw = populations[np.array(fitness).argsort()[-1]]
if args.local_rank == 0:
run_manager.write_log('best cfg: ' + str(best_cfg), 'prune')
run_manager.write_log('best val acc: ' + str(best_acc), 'prune')
fitness = best_acc
cfg = list(best_cfg)
# final fine-tuning
if args.model == 'resnet18':
run_manager.reset_model(
ResNet_ImageNet(num_classes=1000, cfg=cfg, depth=18),
net_origin.cpu())
elif args.model == 'resnet34':
run_manager.reset_model(
ResNet_ImageNet(num_classes=1000, cfg=cfg, depth=34),
net_origin.cpu())
elif args.model == 'resnet50':
run_manager.reset_model(
ResNet_ImageNet(num_classes=1000, cfg=cfg, depth=50),
net_origin.cpu())
elif args.model == 'mobilenet':
run_manager.reset_model(MobileNet(num_classes=1000, cfg=cfg),
net_origin.cpu())
elif args.model == 'mobilenetv2':
run_manager.reset_model(MobileNetV2(num_classes=1000, cfg=cfg),
net_origin.cpu())
