# prepare network net_config_path = '%s/net.config' % args.path if os.path.isfile(net_config_path): # load net from file from models import get_net_by_name net_config = json.load(open(net_config_path, 'r')) net = get_net_by_name(net_config['name']).build_from_config(net_config) else: # build net from args if 'proxyless' in args.net: from models.normal_nets.proxyless_nets import proxyless_base net_config_url = 'https://hanlab.mit.edu/files/proxylessNAS/%s.config' % args.net net = proxyless_base( net_config=net_config_url, n_classes=run_config.data_provider.n_classes, bn_param=(args.bn_momentum, args.bn_eps), dropout_rate=args.dropout, ) else: raise ValueError('do not support: %s' % args.net) # build run manager run_manager = RunManager(args.path, net, run_config, measure_latency=args.latency) run_manager.save_config(print_info=True) # load checkpoints best_model_path = '%s/checkpoint/model_best.pth.tar' % args.path if os.path.isfile(best_model_path):
def main(args, myargs): torch.manual_seed(args.manual_seed) torch.cuda.manual_seed_all(args.manual_seed) np.random.seed(args.manual_seed) # os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu os.makedirs(args.path, exist_ok=True) # prepare run config run_config_path = '%s/run.config' % args.path if os.path.isfile(run_config_path): # load run config from file run_config = json.load(open(run_config_path, 'r')) run_config = ImagenetRunConfig(**run_config) if args.valid_size: run_config.valid_size = args.valid_size else: # build run config from args args.lr_schedule_param = None args.opt_param = { 'momentum': args.momentum, 'nesterov': not args.no_nesterov, } if args.no_decay_keys == 'None': args.no_decay_keys = None run_config = ImagenetRunConfig(**args.__dict__) print('Run config:') for k, v in run_config.config.items(): print('\t%s: %s' % (k, v)) # prepare network net_config_path = '%s/net.config' % args.path if os.path.isfile(net_config_path): # load net from file from models import get_net_by_name net_config = json.load(open(net_config_path, 'r')) net = get_net_by_name(net_config['name']).build_from_config(net_config) else: # build net from args if 'proxyless' in args.net: from models.normal_nets.proxyless_nets import proxyless_base net_config_url = 'https://hanlab.mit.edu/files/proxylessNAS/%s.config' % args.net net = proxyless_base( net_config=net_config_url, n_classes=run_config.data_provider.n_classes, bn_param=(args.bn_momentum, args.bn_eps), dropout_rate=args.dropout, ) else: raise ValueError('do not support: %s' % args.net) # build run manager run_manager = RunManager(args.path, net, run_config, measure_latency=args.latency) run_manager.save_config(print_info=True) # load checkpoints init_path = '%s/init' % args.path if args.resume: run_manager.load_model() if args.train and run_manager.best_acc == 0: loss, acc1, acc5 = run_manager.validate(is_test=False, return_top5=True) run_manager.best_acc = acc1 elif os.path.isfile(init_path): if torch.cuda.is_available(): checkpoint = torch.load(init_path) else: checkpoint = torch.load(init_path, map_location='cpu') if 'state_dict' in checkpoint: checkpoint = checkpoint['state_dict'] run_manager.net.module.load_state_dict(checkpoint) elif 'proxyless' in args.net and not args.train: from utils.latency_estimator import download_url pretrained_weight_url = 'https://hanlab.mit.edu/files/proxylessNAS/%s.pth' % args.net print('Load pretrained weights from %s' % pretrained_weight_url) init_path = download_url(pretrained_weight_url) init = torch.load(init_path, map_location='cpu') net.load_state_dict(init['state_dict']) else: print('Random initialization') # train if args.train: print('Start training') run_manager.train(print_top5=True) run_manager.save_model() output_dict = {} # validate if run_config.valid_size: print('Test on validation set') loss, acc1, acc5 = run_manager.validate(is_test=False, return_top5=True) log = 'valid_loss: %f\t valid_acc1: %f\t valid_acc5: %f' % (loss, acc1, acc5) run_manager.write_log(log, prefix='valid') output_dict = { **output_dict, 'valid_loss': ' % f' % loss, 'valid_acc1': ' % f' % acc1, 'valid_acc5': ' % f' % acc5, 'valid_size': run_config.valid_size } # test print('Test on test set') loss, acc1, acc5 = run_manager.validate(is_test=True, return_top5=True) log = 'test_loss: %f\t test_acc1: %f\t test_acc5: %f' % (loss, acc1, acc5) run_manager.write_log(log, prefix='test') output_dict = { **output_dict, 'test_loss': '%f' % loss, 'test_acc1': '%f' % acc1, 'test_acc5': '%f' % acc5 } json.dump(output_dict, open('%s/output' % args.path, 'w'), indent=4)