plot_shape = (25, 5)
num_plots = plot_shape[0]*plot_shape[1]
fig, axes = plt.subplots(*plot_shape, sharex=True, figsize=(9, 9) )
plt_cts = [0 for i in range(plot_shape[1])]
api = API(args.api_loc)
archs = list(range(ARCH_START, ARCH_END))
colours = ['#811F41', '#A92941', '#D15141', '#EF7941', '#F99C4B']
strs = []
random.shuffle(archs)
for arch in archs:
try:
config = api.get_net_config(arch, 'cifar10')
archinfo = api.query_meta_info_by_index(arch)
acc = archinfo.get_metrics('cifar10-valid', 'x-valid')['accuracy']
network = get_cell_based_tiny_net(config)
network = network.to(device)
jacobs, labels = get_batch_jacobian(network, train_loader, device)
boundaries = [60., 70., 80., 90.]
can_plt, row, col, accrange = decide_plot(acc, plt_cts, plot_shape[0], boundaries)
if not can_plt:
continue
axes[row, col].axis('off')
plot_hist(jacobs, axes[row, col], colours[col])
if row == 0:
axes[row, col].set_title(f'{accrange}')
def main(xargs):
assert torch.cuda.is_available(), 'CUDA is not available.'
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
torch.set_num_threads(xargs.workers)
prepare_seed(xargs.rand_seed)
logger = prepare_logger(args)
train_data, valid_data, xshape, class_num = get_datasets(
xargs.dataset, xargs.data_path, -1)
#config_path = 'configs/nas-benchmark/algos/GDAS.config'
config = load_config(xargs.config_path, {
'class_num': class_num,
'xshape': xshape
}, logger)
search_loader, _, valid_loader = get_nas_search_loaders(
train_data, valid_data, xargs.dataset, 'configs/nas-benchmark/',
config.batch_size, xargs.workers)
logger.log(
'||||||| {:10s} ||||||| Search-Loader-Num={:}, batch size={:}'.format(
xargs.dataset, len(search_loader), config.batch_size))
logger.log('||||||| {:10s} ||||||| Config={:}'.format(
xargs.dataset, config))
search_space = get_search_spaces('cell', xargs.search_space_name)
if xargs.model_config is None and not args.constrain:
model_config = dict2config(
{
'name': 'ProxylessNAS',
'C': xargs.channel,
'N': xargs.num_cells,
'max_nodes': xargs.max_nodes,
'num_classes': class_num,
'space': search_space,
'inp_size': 0,
'affine': False,
'track_running_stats': bool(xargs.track_running_stats)
}, None)
elif xargs.model_config is None:
model_config = dict2config(
{
'name': 'ProxylessNAS',
'C': xargs.channel,
'N': xargs.num_cells,
'max_nodes': xargs.max_nodes,
'num_classes': class_num,
'space': search_space,
'inp_size': 32,
'affine': False,
'track_running_stats': bool(xargs.track_running_stats)
}, None)
else:
model_config = load_config(
xargs.model_config, {
'num_classes': class_num,
'space': search_space,
'affine': False,
'track_running_stats': bool(xargs.track_running_stats)
}, None)
search_model = get_cell_based_tiny_net(model_config)
#logger.log('search-model :\n{:}'.format(search_model))
logger.log('model-config : {:}'.format(model_config))
w_optimizer, w_scheduler, criterion = get_optim_scheduler(
search_model.get_weights(), config)
a_optimizer = torch.optim.Adam(search_model.get_alphas(),
lr=xargs.arch_learning_rate,
betas=(0.5, 0.999),
weight_decay=xargs.arch_weight_decay)
logger.log('w-optimizer : {:}'.format(w_optimizer))
logger.log('a-optimizer : {:}'.format(a_optimizer))
logger.log('w-scheduler : {:}'.format(w_scheduler))
logger.log('criterion : {:}'.format(criterion))
flop, param = get_model_infos(search_model, xshape)
#logger.log('{:}'.format(search_model))
logger.log('FLOP = {:.2f} M, Params = {:.2f} MB'.format(flop, param))
logger.log('search-space [{:} ops] : {:}'.format(len(search_space),
search_space))
if xargs.arch_nas_dataset is None:
api = None
else:
api = API(xargs.arch_nas_dataset)
logger.log('{:} create API = {:} done'.format(time_string(), api))
last_info, model_base_path, model_best_path = logger.path(
'info'), logger.path('model'), logger.path('best')
network, criterion = torch.nn.DataParallel(
search_model).cuda(), criterion.cuda()
#network, criterion = search_model.cuda(), criterion.cuda()
if last_info.exists(): # automatically resume from previous checkpoint
logger.log("=> loading checkpoint of the last-info '{:}' start".format(
last_info))
last_info = torch.load(last_info)
start_epoch = last_info['epoch']
checkpoint = torch.load(last_info['last_checkpoint'])
genotypes = checkpoint['genotypes']
valid_accuracies = checkpoint['valid_accuracies']
search_model.load_state_dict(checkpoint['search_model'])
w_scheduler.load_state_dict(checkpoint['w_scheduler'])
w_optimizer.load_state_dict(checkpoint['w_optimizer'])
a_optimizer.load_state_dict(checkpoint['a_optimizer'])
logger.log(
"=> loading checkpoint of the last-info '{:}' start with {:}-th epoch."
.format(last_info, start_epoch))
else:
logger.log("=> do not find the last-info file : {:}".format(last_info))
start_epoch, valid_accuracies, genotypes = 0, {
'best': -1
}, {
-1: search_model.genotype()
}
# start training
start_time, search_time, epoch_time, total_epoch = time.time(
), AverageMeter(), AverageMeter(), config.epochs + config.warmup
sampled_weights = []
for epoch in range(start_epoch, total_epoch + config.t_epochs):
w_scheduler.update(epoch, 0.0)
need_time = 'Time Left: {:}'.format(
convert_secs2time(
epoch_time.val * (total_epoch - epoch + config.t_epochs),
True))
epoch_str = '{:03d}-{:03d}'.format(epoch, total_epoch)
search_model.set_tau(xargs.tau_max -
(xargs.tau_max - xargs.tau_min) * epoch /
(total_epoch - 1))
logger.log('\n[Search the {:}-th epoch] {:}, tau={:}, LR={:}'.format(
epoch_str, need_time, search_model.get_tau(),
min(w_scheduler.get_lr())))
if epoch < total_epoch:
search_w_loss, search_w_top1, search_w_top5, valid_a_loss , valid_a_top1 , valid_a_top5 \
= search_func(search_loader, network, criterion, w_scheduler, w_optimizer, a_optimizer, epoch_str, xargs.print_freq, logger, xargs.bilevel)
else:
try:
search_w_loss, search_w_top1, search_w_top5, valid_a_loss , valid_a_top1 , valid_a_top5, arch_iter \
= train_func(search_loader, network, criterion, w_scheduler, w_optimizer, epoch_str, xargs.print_freq, sampled_weights[0], arch_iter, logger)
except IndexError:
weights = search_model.sample_weights(100)
sampled_weights.append(weights)
arch_iter = iter(weights)
search_w_loss, search_w_top1, search_w_top5, valid_a_loss , valid_a_top1 , valid_a_top5, arch_iter \
= train_func(search_loader, network, criterion, w_scheduler, w_optimizer, epoch_str, xargs.print_freq, sampled_weights[0], arch_iter, logger)
search_time.update(time.time() - start_time)
logger.log(
'[{:}] searching : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%, time-cost={:.1f} s'
.format(epoch_str, search_w_loss, search_w_top1, search_w_top5,
search_time.sum))
logger.log(
'[{:}] evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'
.format(epoch_str, valid_a_loss, valid_a_top1, valid_a_top5))
if (epoch + 1) % 50 == 0 and not config.t_epochs:
weights = search_model.sample_weights(100)
sampled_weights.append(weights)
elif (epoch + 1) == total_epoch and config.t_epochs:
weights = search_model.sample_weights(100)
sampled_weights.append(weights)
arch_iter = iter(weights)
# validate with single arch
single_weight = search_model.sample_weights(1)[0]
single_valid_acc = AverageMeter()
network.eval()
for i in range(10):
try:
val_input, val_target = next(valid_iter)
except Exception as e:
valid_iter = iter(valid_loader)
val_input, val_target = next(valid_iter)
n_val = val_input.size(0)
with torch.no_grad():
val_target = val_target.cuda(non_blocking=True)
_, logits, _ = network(val_input, weights=single_weight)
val_acc1, val_acc5 = obtain_accuracy(logits.data,
val_target.data,
topk=(1, 5))
single_valid_acc.update(val_acc1.item(), n_val)
logger.log('[{:}] valid : accuracy = {:.2f}'.format(
epoch_str, single_valid_acc.avg))
# check the best accuracy
valid_accuracies[epoch] = valid_a_top1
if valid_a_top1 > valid_accuracies['best']:
valid_accuracies['best'] = valid_a_top1
genotypes['best'] = search_model.genotype()
find_best = True
else:
find_best = False
if epoch < total_epoch:
genotypes[epoch] = search_model.genotype()
logger.log('<<<--->>> The {:}-th epoch : {:}'.format(
epoch_str, genotypes[epoch]))
# save checkpoint
save_path = save_checkpoint(
{
'epoch': epoch + 1,
'args': deepcopy(xargs),
'search_model': search_model.state_dict(),
'w_optimizer': w_optimizer.state_dict(),
'a_optimizer': a_optimizer.state_dict(),
'w_scheduler': w_scheduler.state_dict(),
'genotypes': genotypes,
'valid_accuracies': valid_accuracies
}, model_base_path, logger)
last_info = save_checkpoint(
{
'epoch': epoch + 1,
'args': deepcopy(args),
'last_checkpoint': save_path,
}, logger.path('info'), logger)
if find_best:
logger.log(
'<<<--->>> The {:}-th epoch : find the highest validation accuracy : {:.2f}%.'
.format(epoch_str, valid_a_top1))
copy_checkpoint(model_base_path, model_best_path, logger)
with torch.no_grad():
logger.log('{:}'.format(search_model.show_alphas()))
if api is not None and epoch < total_epoch:
logger.log('{:}'.format(api.query_by_arch(genotypes[epoch])))
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
network.eval()
# Evaluate the architectures sampled throughout the search
for i in range(len(sampled_weights) - 1):
logger.log('Sample eval : epoch {}'.format((i + 1) * 50 - 1))
for w in sampled_weights[i]:
sample_valid_acc = AverageMeter()
for i in range(10):
try:
val_input, val_target = next(valid_iter)
except Exception as e:
valid_iter = iter(valid_loader)
val_input, val_target = next(valid_iter)
n_val = val_input.size(0)
with torch.no_grad():
val_target = val_target.cuda(non_blocking=True)
_, logits, _ = network(val_input, weights=w)
val_acc1, val_acc5 = obtain_accuracy(logits.data,
val_target.data,
topk=(1, 5))
sample_valid_acc.update(val_acc1.item(), n_val)
w_gene = search_model.genotype(w)
if api is not None:
ind = api.query_index_by_arch(w_gene)
info = api.query_meta_info_by_index(ind)
metrics = info.get_metrics('cifar10', 'ori-test')
acc = metrics['accuracy']
else:
acc = 0.0
logger.log(
'sample valid : val_acc = {:.2f} test_acc = {:.2f}'.format(
sample_valid_acc.avg, acc))
# Evaluate the final sampling separately to find the top 10 architectures
logger.log('Final sample eval')
final_archs = []
for w in sampled_weights[-1]:
sample_valid_acc = AverageMeter()
for i in range(10):
try:
val_input, val_target = next(valid_iter)
except Exception as e:
valid_iter = iter(valid_loader)
val_input, val_target = next(valid_iter)
n_val = val_input.size(0)
with torch.no_grad():
val_target = val_target.cuda(non_blocking=True)
_, logits, _ = network(val_input, weights=w)
val_acc1, val_acc5 = obtain_accuracy(logits.data,
val_target.data,
topk=(1, 5))
sample_valid_acc.update(val_acc1.item(), n_val)
w_gene = search_model.genotype(w)
if api is not None:
ind = api.query_index_by_arch(w_gene)
info = api.query_meta_info_by_index(ind)
metrics = info.get_metrics('cifar10', 'ori-test')
acc = metrics['accuracy']
else:
acc = 0.0
logger.log('sample valid : val_acc = {:.2f} test_acc = {:.2f}'.format(
sample_valid_acc.avg, acc))
final_archs.append((w, sample_valid_acc.avg))
top_10 = sorted(final_archs, key=lambda x: x[1], reverse=True)[:10]
# Evaluate the top 10 architectures on the entire validation set
logger.log('Evaluating top archs')
for w, prev_acc in top_10:
full_valid_acc = AverageMeter()
for val_input, val_target in valid_loader:
n_val = val_input.size(0)
with torch.no_grad():
val_target = val_target.cuda(non_blocking=True)
_, logits, _ = network(val_input, weights=w)
val_acc1, val_acc5 = obtain_accuracy(logits.data,
val_target.data,
topk=(1, 5))
full_valid_acc.update(val_acc1.item(), n_val)
w_gene = search_model.genotype(w)
logger.log('genotype {}'.format(w_gene))
if api is not None:
ind = api.query_index_by_arch(w_gene)
info = api.query_meta_info_by_index(ind)
metrics = info.get_metrics('cifar10', 'ori-test')
acc = metrics['accuracy']
else:
acc = 0.0
logger.log(
'full valid : val_acc = {:.2f} test_acc = {:.2f} pval_acc = {:.2f}'
.format(full_valid_acc.avg, acc, prev_acc))
logger.log('\n' + '-' * 100)
# check the performance from the architecture dataset
logger.log(
'ProxylessNAS : run {:} epochs, cost {:.1f} s, last-geno is {:}.'.
format(total_epoch, search_time.sum, genotypes[total_epoch - 1]))
if api is not None:
logger.log('{:}'.format(api.query_by_arch(genotypes[total_epoch - 1])))
logger.close()
