best_cfg = None
best_cfg_raw = None
if args.local_rank == 0:
print('=' * 40 + 'initialization' + '=' * 40)
cfg_len = len(net.config['cfg_base'])
for p in range(n_population):
rand_sample = np.random.rand(cfg_len)
rand_cfg = raw2cfg(net,
np.array(rand_sample),
target_flops,
div=args.div)
populations.append(np.array(rand_sample))
cfgs.append(rand_cfg)
fitness.append(cfg2fitness(rand_cfg))
if args.local_rank == 0:
run_manager.write_log('cfgs: ' + str(cfgs), 'valid')
run_manager.write_log('fitness: ' + str(fitness), 'valid')
run_manager.write_log('avg fitness: ' + str(np.mean(fitness)), 'prune')
run_manager.write_log('best fitness: ' + str(np.max(fitness)), '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')
# evolutionary search
if args.local_rank == 0:
print('=' * 40 + 'evolving' + '=' * 40)
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:
print('Random initialization')
# train
if args.train:
print('Start training: %d' % args.local_rank)
if args.local_rank == 0:
print(net)
run_manager.train(print_top5=True)
if args.local_rank == 0:
run_manager.save_model()
output_dict = {}
# 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)
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)
