if args.valid_size:
run_config.valid_size = args.valid_size
else:
raise FileNotFoundError
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:
raise FileNotFoundError
# load checkpoints
best_model_path = '%s/checkpoint/model_best.pth.tar' % args.path
if os.path.isfile(best_model_path):
print("Use trained model parameters from model_best.pth.tar")
init_path = best_model_path
else:
print("Fallback to using parameters from init file")
init_path = '%s/init' % args.path
if not os.path.isfile(init_path):
raise FileNotFoundError
if torch.cuda.is_available():
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)