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)
if xargs.dataset == 'cifar10' or xargs.dataset == 'cifar100':
split_Fpath = 'configs/nas-benchmark/cifar-split.txt'
cifar_split = load_config(split_Fpath, None, None)
train_split, valid_split = cifar_split.train, cifar_split.valid
logger.log('Load split file from {:}'.format(split_Fpath))
elif xargs.dataset.startswith('ImageNet16'):
split_Fpath = 'configs/nas-benchmark/{:}-split.txt'.format(
xargs.dataset)
imagenet16_split = load_config(split_Fpath, None, None)
train_split, valid_split = imagenet16_split.train, imagenet16_split.valid
logger.log('Load split file from {:}'.format(split_Fpath))
else:
raise ValueError('invalid dataset : {:}'.format(xargs.dataset))
config_path = 'configs/nas-benchmark/algos/DARTS.config'
config = load_config(config_path, {
'class_num': class_num,
'xshape': xshape
}, logger)
# To split data
train_data_v2 = deepcopy(train_data)
train_data_v2.transform = valid_data.transform
valid_data = train_data_v2
search_data = SearchDataset(xargs.dataset, train_data, train_split,
valid_split)
# data loader
search_loader = torch.utils.data.DataLoader(search_data,
batch_size=config.batch_size,
shuffle=True,
num_workers=xargs.workers,
pin_memory=True)
valid_loader = torch.utils.data.DataLoader(
valid_data,
batch_size=config.batch_size,
sampler=torch.utils.data.sampler.SubsetRandomSampler(valid_split),
num_workers=xargs.workers,
pin_memory=True)
logger.log(
'||||||| {:10s} ||||||| Search-Loader-Num={:}, Valid-Loader-Num={:}, batch size={:}'
.format(xargs.dataset, len(search_loader), len(valid_loader),
config.batch_size))
logger.log('||||||| {:10s} ||||||| Config={:}'.format(
xargs.dataset, config))
search_space = get_search_spaces('cell', xargs.search_space_name)
model_config = dict2config(
{
'name': 'DARTS-V1',
'C': xargs.channel,
'N': xargs.num_cells,
'max_nodes': xargs.max_nodes,
'num_classes': class_num,
'space': search_space
}, None)
search_model = get_cell_based_tiny_net(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))
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()
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}, {}
# start training
start_time, epoch_time, total_epoch = time.time(), AverageMeter(
), config.epochs + config.warmup
for epoch in range(start_epoch, total_epoch):
w_scheduler.update(epoch, 0.0)
need_time = 'Time Left: {:}'.format(
convert_secs2time(epoch_time.val * (total_epoch - epoch), True))
epoch_str = '{:03d}-{:03d}'.format(epoch, total_epoch)
logger.log('\n[Search the {:}-th epoch] {:}, LR={:}'.format(
epoch_str, need_time, min(w_scheduler.get_lr())))
search_w_loss, search_w_top1, search_w_top5 = search_func(
search_loader, network, criterion, w_scheduler, w_optimizer,
a_optimizer, epoch_str, xargs.print_freq, logger)
logger.log(
'[{:}] searching : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'
.format(epoch_str, search_w_loss, search_w_top1, search_w_top5))
valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
valid_loader, network, criterion)
logger.log(
'[{:}] evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'
.format(epoch_str, valid_a_loss, valid_a_top1, valid_a_top5))
# 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
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('arch-parameters :\n{:}'.format(
nn.functional.softmax(search_model.arch_parameters,
dim=-1).cpu()))
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log('\n' + '-' * 100)
# check the performance from the architecture dataset
#if xargs.arch_nas_dataset is None or not os.path.isfile(xargs.arch_nas_dataset):
# logger.log('Can not find the architecture dataset : {:}.'.format(xargs.arch_nas_dataset))
#else:
# nas_bench = NASBenchmarkAPI(xargs.arch_nas_dataset)
# geno = genotypes[total_epoch-1]
# logger.log('The last model is {:}'.format(geno))
# info = nas_bench.query_by_arch( geno )
# if info is None: logger.log('Did not find this architecture : {:}.'.format(geno))
# else : logger.log('{:}'.format(info))
# logger.log('-'*100)
# geno = genotypes['best']
# logger.log('The best model is {:}'.format(geno))
# info = nas_bench.query_by_arch( geno )
# if info is None: logger.log('Did not find this architecture : {:}.'.format(geno))
# else : logger.log('{:}'.format(info))
logger.close()
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)
if os.path.isdir(xargs.save_dir):
if click.confirm(
'\nSave directory already exists in {}. Erase?'.format(
xargs.save_dir, default=False)):
os.system('rm -r ' + xargs.save_dir)
assert not os.path.exists(xargs.save_dir)
os.mkdir(xargs.save_dir)
logger = prepare_logger(args)
writer = SummaryWriter(xargs.save_dir)
perturb_alpha = None
if xargs.perturb:
perturb_alpha = random_alpha
train_data, valid_data, xshape, class_num = get_datasets(
xargs.dataset, xargs.data_path, -1)
# config_path = 'configs/nas-benchmark/algos/DARTS.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={:}, Valid-Loader-Num={:}, batch size={:}'
.format(xargs.dataset, len(search_loader), len(valid_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:
model_config = dict2config(
{
'name': xargs.model,
'C': xargs.channel,
'N': xargs.num_cells,
'max_nodes': xargs.max_nodes,
'num_classes': class_num,
'space': search_space,
'affine': bool(xargs.affine),
'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': bool(xargs.affine),
'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))
w_optimizer, w_scheduler, criterion = get_optim_scheduler(
search_model.get_weights(), config, xargs.weight_learning_rate)
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))
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()
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
start_time, search_time, epoch_time = time.time(), AverageMeter(
), AverageMeter()
total_epoch = config.epochs + config.warmup
assert 0 < xargs.early_stop_epoch <= total_epoch - 1
for epoch in range(start_epoch, total_epoch):
if epoch >= xargs.early_stop_epoch:
logger.log(f"Early stop @ {epoch} epoch.")
break
if xargs.perturb:
epsilon_alpha = 0.03 + (xargs.epsilon_alpha -
0.03) * epoch / total_epoch
logger.log(f'epoch {epoch} epsilon_alpha {epsilon_alpha}')
else:
epsilon_alpha = None
w_scheduler.update(epoch, 0.0)
need_time = 'Time Left: {:}'.format(
convert_secs2time(epoch_time.val * (total_epoch - epoch), True))
epoch_str = '{:03d}-{:03d}'.format(epoch, total_epoch)
logger.log('\n[Search the {:}-th epoch] {:}, LR={:}'.format(
epoch_str, need_time, min(w_scheduler.get_lr())))
search_w_loss, search_w_top1, search_w_top5, search_a_loss, search_a_top1, search_a_top5 = search_func(
search_loader, network, criterion, w_scheduler, w_optimizer,
a_optimizer, epoch_str, xargs.print_freq, logger,
xargs.gradient_clip, perturb_alpha, epsilon_alpha)
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))
valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
valid_loader, network, criterion)
writer.add_scalar('search/weight_loss', search_w_loss, epoch)
writer.add_scalar('search/weight_top1_acc', search_w_top1, epoch)
writer.add_scalar('search/weight_top5_acc', search_w_top5, epoch)
writer.add_scalar('search/arch_loss', search_a_loss, epoch)
writer.add_scalar('search/arch_top1_acc', search_a_top1, epoch)
writer.add_scalar('search/arch_top5_acc', search_a_top5, epoch)
writer.add_scalar('evaluate/loss', valid_a_loss, epoch)
writer.add_scalar('evaluate/top1_acc', valid_a_top1, epoch)
writer.add_scalar('evaluate/top5_acc', valid_a_top5, epoch)
logger.log(
'[{:}] evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'
.format(epoch_str, valid_a_loss, valid_a_top1, valid_a_top5))
writer.add_scalar('entropy', search_model.entropy, epoch)
per_edge_dict = get_per_egde_value_dict(search_model.arch_parameters)
for edge_name, edge_val in per_edge_dict.items():
writer.add_scalars(f"cell/{edge_name}", edge_val, epoch)
# 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
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)
save_checkpoint(
{
'epoch': epoch + 1,
'args': deepcopy(args),
'last_checkpoint': save_path,
}, logger.path('info'), logger)
if xargs.snapshoot > 0 and epoch % xargs.snapshoot == 0:
save_checkpoint(
{
'epoch': epoch + 1,
'args': deepcopy(args),
'search_model': search_model.state_dict(),
},
os.path.join(str(logger.model_dir),
f"checkpoint_epoch{epoch}.pth"), 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:
logger.log('{:}'.format(api.query_by_arch(genotypes[epoch])))
index = api.query_index_by_arch(genotypes[epoch])
info = api.query_meta_info_by_index(
index) # This is an instance of `ArchResults`
res_metrics = info.get_metrics(
f'{xargs.dataset}',
'ori-test') # This is a dict with metric names as keys
# cost_metrics = info.get_comput_costs('cifar10')
writer.add_scalar(f'{xargs.dataset}_ground_acc_ori-test',
res_metrics['accuracy'], epoch)
writer.add_scalar(f'{xargs.dataset}_search_acc', valid_a_top1,
epoch)
if xargs.dataset.lower() != 'cifar10':
writer.add_scalar(
f'{xargs.dataset}_ground_acc_x-test',
info.get_metrics(f'{xargs.dataset}', 'x-test')['accuracy'],
epoch)
if find_best:
valid_accuracies['best_gt'] = res_metrics['accuracy']
writer.add_scalar(f"{xargs.dataset}_cur_best_gt_acc_ori-test",
valid_accuracies['best_gt'], epoch)
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log('\n' + '-' * 100)
logger.log('{:} : run {:} epochs, cost {:.1f} s, last-geno is {:}.'.format(
args.model, xargs.early_stop_epoch, search_time.sum,
genotypes[xargs.early_stop_epoch - 1]))
if api is not None:
logger.log('{:}'.format(
api.query_by_arch(genotypes[xargs.early_stop_epoch - 1])))
logger.close()
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, test_data, xshape, class_num = get_datasets(
xargs.dataset, xargs.data_path, -1)
logger.log('use config from : {:}'.format(xargs.config_path))
config = load_config(xargs.config_path, {
'class_num': class_num,
'xshape': xshape
}, logger)
_, train_loader, valid_loader = get_nas_search_loaders(
train_data, test_data, xargs.dataset, 'configs/nas-benchmark/',
config.batch_size, xargs.workers)
# since ENAS will train the controller on valid-loader, we need to use train transformation for valid-loader
valid_loader.dataset.transform = deepcopy(train_loader.dataset.transform)
if hasattr(valid_loader.dataset, 'transforms'):
valid_loader.dataset.transforms = deepcopy(
train_loader.dataset.transforms)
# data loader
logger.log(
'||||||| {:10s} ||||||| Train-Loader-Num={:}, Valid-Loader-Num={:}, batch size={:}'
.format(xargs.dataset, len(train_loader), len(valid_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:
model_config = dict2config(
{
'name': 'ENAS',
'C': xargs.channel,
'N': xargs.num_cells,
'max_nodes': xargs.max_nodes,
'num_classes': class_num,
'space': search_space,
'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)
shared_cnn = get_cell_based_tiny_net(model_config)
controller = shared_cnn.create_controller()
w_optimizer, w_scheduler, criterion = get_optim_scheduler(
shared_cnn.parameters(), config)
a_optimizer = torch.optim.Adam(controller.parameters(),
lr=config.controller_lr,
betas=config.controller_betas,
eps=config.controller_eps)
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(shared_cnn, xshape)
#logger.log('{:}'.format(shared_cnn))
#logger.log('FLOP = {:.2f} M, Params = {:.2f} MB'.format(flop, param))
logger.log('search-space : {:}'.format(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))
shared_cnn, controller, criterion = torch.nn.DataParallel(
shared_cnn).cuda(), controller.cuda(), criterion.cuda()
last_info, model_base_path, model_best_path = logger.path(
'info'), logger.path('model'), logger.path('best')
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']
baseline = checkpoint['baseline']
valid_accuracies = checkpoint['valid_accuracies']
shared_cnn.load_state_dict(checkpoint['shared_cnn'])
controller.load_state_dict(checkpoint['controller'])
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, baseline = 0, {
'best': -1
}, {}, None
# start training
start_time, search_time, epoch_time, total_epoch = time.time(
), AverageMeter(), AverageMeter(), config.epochs + config.warmup
for epoch in range(start_epoch, total_epoch):
w_scheduler.update(epoch, 0.0)
need_time = 'Time Left: {:}'.format(
convert_secs2time(epoch_time.val * (total_epoch - epoch), True))
epoch_str = '{:03d}-{:03d}'.format(epoch, total_epoch)
logger.log(
'\n[Search the {:}-th epoch] {:}, LR={:}, baseline={:}'.format(
epoch_str, need_time, min(w_scheduler.get_lr()), baseline))
cnn_loss, cnn_top1, cnn_top5 = train_shared_cnn(
train_loader, shared_cnn, controller, criterion, w_scheduler,
w_optimizer, epoch_str, xargs.print_freq, logger)
logger.log(
'[{:}] shared-cnn : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'
.format(epoch_str, cnn_loss, cnn_top1, cnn_top5))
ctl_loss, ctl_acc, ctl_baseline, ctl_reward, baseline \
= train_controller(valid_loader, shared_cnn, controller, criterion, a_optimizer, \
dict2config({'baseline': baseline,
'ctl_train_steps': xargs.controller_train_steps, 'ctl_num_aggre': xargs.controller_num_aggregate,
'ctl_entropy_w': xargs.controller_entropy_weight,
'ctl_bl_dec' : xargs.controller_bl_dec}, None), \
epoch_str, xargs.print_freq, logger)
search_time.update(time.time() - start_time)
logger.log(
'[{:}] controller : loss={:.2f}, accuracy={:.2f}%, baseline={:.2f}, reward={:.2f}, current-baseline={:.4f}, time-cost={:.1f} s'
.format(epoch_str, ctl_loss, ctl_acc, ctl_baseline, ctl_reward,
baseline, search_time.sum))
best_arch, _ = get_best_arch(controller, shared_cnn, valid_loader)
shared_cnn.module.update_arch(best_arch)
_, best_valid_acc, _ = valid_func(valid_loader, shared_cnn, criterion)
genotypes[epoch] = best_arch
# check the best accuracy
valid_accuracies[epoch] = best_valid_acc
if best_valid_acc > valid_accuracies['best']:
valid_accuracies['best'] = best_valid_acc
genotypes['best'] = best_arch
find_best = True
else:
find_best = False
logger.log('<<<--->>> The {:}-th epoch : {:}'.format(
epoch_str, genotypes[epoch]))
# save checkpoint
save_path = save_checkpoint(
{
'epoch': epoch + 1,
'args': deepcopy(xargs),
'baseline': baseline,
'shared_cnn': shared_cnn.state_dict(),
'controller': controller.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, best_valid_acc))
copy_checkpoint(model_base_path, model_best_path, logger)
if api is not None:
logger.log('{:}'.format(api.query_by_arch(genotypes[epoch])))
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log('\n' + '-' * 100)
logger.log('During searching, the best architecture is {:}'.format(
genotypes['best']))
logger.log('Its accuracy is {:.2f}%'.format(valid_accuracies['best']))
logger.log('Randomly select {:} architectures and select the best.'.format(
xargs.controller_num_samples))
start_time = time.time()
final_arch, _ = get_best_arch(controller, shared_cnn, valid_loader,
xargs.controller_num_samples)
search_time.update(time.time() - start_time)
shared_cnn.module.update_arch(final_arch)
final_loss, final_top1, final_top5 = valid_func(valid_loader, shared_cnn,
criterion)
logger.log('The Selected Final Architecture : {:}'.format(final_arch))
logger.log('Loss={:.3f}, Accuracy@1={:.2f}%, Accuracy@5={:.2f}%'.format(
final_loss, final_top1, final_top5))
logger.log(
'ENAS : run {:} epochs, cost {:.1f} s, last-geno is {:}.'.format(
total_epoch, search_time.sum, final_arch))
if api is not None: logger.log('{:}'.format(api.query_by_arch(final_arch)))
logger.close()
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 = 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={:}, Valid-Loader-Num={:}, batch size={:}"
.format(xargs.dataset, len(search_loader), len(valid_loader),
config.batch_size))
logger.log("||||||| {:10s} ||||||| Config={:}".format(
xargs.dataset, config))
search_space = get_search_spaces("cell", xargs.search_space_name)
model_config = dict2config(
{
"name": "DARTS-V2",
"C": xargs.channel,
"N": xargs.num_cells,
"max_nodes": xargs.max_nodes,
"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))
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))
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()
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,
)
for epoch in range(start_epoch, total_epoch):
w_scheduler.update(epoch, 0.0)
need_time = "Time Left: {:}".format(
convert_secs2time(epoch_time.val * (total_epoch - epoch), True))
epoch_str = "{:03d}-{:03d}".format(epoch, total_epoch)
min_LR = min(w_scheduler.get_lr())
logger.log("\n[Search the {:}-th epoch] {:}, LR={:}".format(
epoch_str, need_time, min_LR))
search_w_loss, search_w_top1, search_w_top5 = search_func(
search_loader,
network,
criterion,
w_scheduler,
w_optimizer,
a_optimizer,
epoch_str,
xargs.print_freq,
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))
valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
valid_loader, network, criterion)
logger.log(
"[{:}] evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%"
.format(epoch_str, valid_a_loss, valid_a_top1, valid_a_top5))
# 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
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("arch-parameters :\n{:}".format(
nn.functional.softmax(search_model.arch_parameters,
dim=-1).cpu()))
if api is not None:
logger.log("{:}".format(api.query_by_arch(genotypes[epoch],
"200")))
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log("\n" + "-" * 100)
# check the performance from the architecture dataset
logger.log(
"DARTS-V2 : 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]),
"200"))
logger.close()
def main(args):
assert torch.cuda.is_available(), 'CUDA is not available.'
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
#torch.backends.cudnn.deterministic = True
torch.set_num_threads(args.workers)
prepare_seed(args.rand_seed)
logger = prepare_logger(args)
train_data, valid_data, xshape, class_num = get_datasets(
args.dataset, args.data_path, args.cutout_length)
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
valid_loader = torch.utils.data.DataLoader(valid_data,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# get configures
model_config = load_config(args.model_config, {'class_num': class_num},
logger)
optim_config = load_config(
args.optim_config, {
'class_num': class_num,
'KD_alpha': args.KD_alpha,
'KD_temperature': args.KD_temperature
}, logger)
# load checkpoint
teacher_base = load_net_from_checkpoint(args.KD_checkpoint)
teacher = torch.nn.DataParallel(teacher_base).cuda()
base_model = obtain_model(model_config)
flop, param = get_model_infos(base_model, xshape)
logger.log('Student ====>>>>:\n{:}'.format(base_model))
logger.log('Teacher ====>>>>:\n{:}'.format(teacher_base))
logger.log('model information : {:}'.format(base_model.get_message()))
logger.log('-' * 50)
logger.log('Params={:.2f} MB, FLOPs={:.2f} M ... = {:.2f} G'.format(
param, flop, flop / 1e3))
logger.log('-' * 50)
logger.log('train_data : {:}'.format(train_data))
logger.log('valid_data : {:}'.format(valid_data))
optimizer, scheduler, criterion = get_optim_scheduler(
base_model.parameters(), optim_config)
logger.log('optimizer : {:}'.format(optimizer))
logger.log('scheduler : {:}'.format(scheduler))
logger.log('criterion : {:}'.format(criterion))
last_info, model_base_path, model_best_path = logger.path(
'info'), logger.path('model'), logger.path('best')
network, criterion = torch.nn.DataParallel(
base_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'] + 1
checkpoint = torch.load(last_info['last_checkpoint'])
base_model.load_state_dict(checkpoint['base-model'])
scheduler.load_state_dict(checkpoint['scheduler'])
optimizer.load_state_dict(checkpoint['optimizer'])
valid_accuracies = checkpoint['valid_accuracies']
max_bytes = checkpoint['max_bytes']
logger.log(
"=> loading checkpoint of the last-info '{:}' start with {:}-th epoch."
.format(last_info, start_epoch))
elif args.resume is not None:
assert Path(
args.resume).exists(), 'Can not find the resume file : {:}'.format(
args.resume)
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch'] + 1
base_model.load_state_dict(checkpoint['base-model'])
scheduler.load_state_dict(checkpoint['scheduler'])
optimizer.load_state_dict(checkpoint['optimizer'])
valid_accuracies = checkpoint['valid_accuracies']
max_bytes = checkpoint['max_bytes']
logger.log(
"=> loading checkpoint from '{:}' start with {:}-th epoch.".format(
args.resume, start_epoch))
elif args.init_model is not None:
assert Path(args.init_model).exists(
), 'Can not find the initialization file : {:}'.format(args.init_model)
checkpoint = torch.load(args.init_model)
base_model.load_state_dict(checkpoint['base-model'])
start_epoch, valid_accuracies, max_bytes = 0, {'best': -1}, {}
logger.log('=> initialize the model from {:}'.format(args.init_model))
else:
logger.log("=> do not find the last-info file : {:}".format(last_info))
start_epoch, valid_accuracies, max_bytes = 0, {'best': -1}, {}
train_func, valid_func = get_procedures(args.procedure)
total_epoch = optim_config.epochs + optim_config.warmup
# Main Training and Evaluation Loop
start_time = time.time()
epoch_time = AverageMeter()
for epoch in range(start_epoch, total_epoch):
scheduler.update(epoch, 0.0)
need_time = 'Time Left: {:}'.format(
convert_secs2time(epoch_time.avg * (total_epoch - epoch), True))
epoch_str = 'epoch={:03d}/{:03d}'.format(epoch, total_epoch)
LRs = scheduler.get_lr()
find_best = False
logger.log(
'\n***{:s}*** start {:s} {:s}, LR=[{:.6f} ~ {:.6f}], scheduler={:}'
.format(time_string(), epoch_str, need_time, min(LRs), max(LRs),
scheduler))
# train for one epoch
train_loss, train_acc1, train_acc5 = train_func(
train_loader, teacher, network, criterion, scheduler, optimizer,
optim_config, epoch_str, args.print_freq, logger)
# log the results
logger.log(
'***{:s}*** TRAIN [{:}] loss = {:.6f}, accuracy-1 = {:.2f}, accuracy-5 = {:.2f}'
.format(time_string(), epoch_str, train_loss, train_acc1,
train_acc5))
# evaluate the performance
if (epoch % args.eval_frequency == 0) or (epoch + 1 == total_epoch):
logger.log('-' * 150)
valid_loss, valid_acc1, valid_acc5 = valid_func(
valid_loader, teacher, network, criterion, optim_config,
epoch_str, args.print_freq_eval, logger)
valid_accuracies[epoch] = valid_acc1
logger.log(
'***{:s}*** VALID [{:}] loss = {:.6f}, accuracy@1 = {:.2f}, accuracy@5 = {:.2f} | Best-Valid-Acc@1={:.2f}, Error@1={:.2f}'
.format(time_string(), epoch_str, valid_loss, valid_acc1,
valid_acc5, valid_accuracies['best'],
100 - valid_accuracies['best']))
if valid_acc1 > valid_accuracies['best']:
valid_accuracies['best'] = valid_acc1
find_best = True
logger.log(
'Currently, the best validation accuracy found at {:03d}-epoch :: acc@1={:.2f}, acc@5={:.2f}, error@1={:.2f}, error@5={:.2f}, save into {:}.'
.format(epoch, valid_acc1, valid_acc5, 100 - valid_acc1,
100 - valid_acc5, model_best_path))
num_bytes = torch.cuda.max_memory_cached(
next(network.parameters()).device) * 1.0
logger.log(
'[GPU-Memory-Usage on {:} is {:} bytes, {:.2f} KB, {:.2f} MB, {:.2f} GB.]'
.format(
next(network.parameters()).device, int(num_bytes),
num_bytes / 1e3, num_bytes / 1e6, num_bytes / 1e9))
max_bytes[epoch] = num_bytes
if epoch % 10 == 0: torch.cuda.empty_cache()
# save checkpoint
save_path = save_checkpoint(
{
'epoch': epoch,
'args': deepcopy(args),
'max_bytes': deepcopy(max_bytes),
'FLOP': flop,
'PARAM': param,
'valid_accuracies': deepcopy(valid_accuracies),
'model-config': model_config._asdict(),
'optim-config': optim_config._asdict(),
'base-model': base_model.state_dict(),
'scheduler': scheduler.state_dict(),
'optimizer': optimizer.state_dict(),
}, model_base_path, logger)
if find_best: copy_checkpoint(model_base_path, model_best_path, logger)
last_info = save_checkpoint(
{
'epoch': epoch,
'args': deepcopy(args),
'last_checkpoint': save_path,
}, logger.path('info'), logger)
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log('\n' + '-' * 200)
logger.log('||| Params={:.2f} MB, FLOPs={:.2f} M ... = {:.2f} G'.format(
param, flop, flop / 1e3))
logger.log(
'Finish training/validation in {:} with Max-GPU-Memory of {:.2f} MB, and save final checkpoint into {:}'
.format(convert_secs2time(epoch_time.sum, True),
max(v for k, v in max_bytes.items()) / 1e6,
logger.path('info')))
logger.log('-' * 200 + '\n')
logger.close()
def main(args):
assert torch.cuda.is_available(), 'CUDA is not available.'
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
#torch.backends.cudnn.deterministic = True
torch.set_num_threads(args.workers)
prepare_seed(args.rand_seed)
logger = prepare_logger(args)
# prepare dataset
train_data, valid_data, xshape, class_num = get_datasets(
args.dataset, args.data_path, args.cutout_length)
#train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True , num_workers=args.workers, pin_memory=True)
valid_loader = torch.utils.data.DataLoader(valid_data,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
split_file_path = Path(args.split_path)
assert split_file_path.exists(), '{:} does not exist'.format(
split_file_path)
split_info = torch.load(split_file_path)
train_split, valid_split = split_info['train'], split_info['valid']
assert len(
set(train_split).intersection(set(valid_split))
) == 0, 'There should be 0 element that belongs to both train and valid'
assert len(train_split) + len(valid_split) == len(
train_data), '{:} + {:} vs {:}'.format(len(train_split),
len(valid_split),
len(train_data))
search_dataset = SearchDataset(args.dataset, train_data, train_split,
valid_split)
search_train_loader = torch.utils.data.DataLoader(
train_data,
batch_size=args.batch_size,
sampler=torch.utils.data.sampler.SubsetRandomSampler(train_split),
pin_memory=True,
num_workers=args.workers)
search_valid_loader = torch.utils.data.DataLoader(
train_data,
batch_size=args.batch_size,
sampler=torch.utils.data.sampler.SubsetRandomSampler(valid_split),
pin_memory=True,
num_workers=args.workers)
search_loader = torch.utils.data.DataLoader(search_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
sampler=None)
# get configures
if args.ablation_num_select is None or args.ablation_num_select <= 0:
model_config = load_config(args.model_config, {
'class_num': class_num,
'search_mode': 'shape'
}, logger)
else:
model_config = load_config(
args.model_config, {
'class_num': class_num,
'search_mode': 'ablation',
'num_random_select': args.ablation_num_select
}, logger)
# obtain the model
search_model = obtain_search_model(model_config)
MAX_FLOP, param = get_model_infos(search_model, xshape)
optim_config = load_config(args.optim_config, {
'class_num': class_num,
'FLOP': MAX_FLOP
}, logger)
logger.log('Model Information : {:}'.format(search_model.get_message()))
logger.log('MAX_FLOP = {:} M'.format(MAX_FLOP))
logger.log('Params = {:} M'.format(param))
logger.log('train_data : {:}'.format(train_data))
logger.log('search-data: {:}'.format(search_dataset))
logger.log('search_train_loader : {:} samples'.format(len(train_split)))
logger.log('search_valid_loader : {:} samples'.format(len(valid_split)))
base_optimizer, scheduler, criterion = get_optim_scheduler(
search_model.base_parameters(), optim_config)
arch_optimizer = torch.optim.Adam(search_model.arch_parameters(
optim_config.arch_LR),
lr=optim_config.arch_LR,
betas=(0.5, 0.999),
weight_decay=optim_config.arch_decay)
logger.log('base-optimizer : {:}'.format(base_optimizer))
logger.log('arch-optimizer : {:}'.format(arch_optimizer))
logger.log('scheduler : {:}'.format(scheduler))
logger.log('criterion : {:}'.format(criterion))
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()
# load checkpoint
if last_info.exists() or (args.resume is not None and osp.isfile(
args.resume)): # automatically resume from previous checkpoint
if args.resume is not None and osp.isfile(args.resume):
resume_path = Path(args.resume)
elif last_info.exists():
resume_path = last_info
else:
raise ValueError('Something is wrong.')
logger.log("=> loading checkpoint of the last-info '{:}' start".format(
resume_path))
checkpoint = torch.load(resume_path)
if 'last_checkpoint' in checkpoint:
last_checkpoint_path = checkpoint['last_checkpoint']
if not last_checkpoint_path.exists():
logger.log('Does not find {:}, try another path'.format(
last_checkpoint_path))
last_checkpoint_path = resume_path.parent / last_checkpoint_path.parent.name / last_checkpoint_path.name
assert last_checkpoint_path.exists(
), 'can not find the checkpoint from {:}'.format(
last_checkpoint_path)
checkpoint = torch.load(last_checkpoint_path)
start_epoch = checkpoint['epoch'] + 1
#for key, value in checkpoint['search_model'].items():
# print('K {:} = Shape={:}'.format(key, value.shape))
search_model.load_state_dict(checkpoint['search_model'])
scheduler.load_state_dict(checkpoint['scheduler'])
base_optimizer.load_state_dict(checkpoint['base_optimizer'])
arch_optimizer.load_state_dict(checkpoint['arch_optimizer'])
valid_accuracies = checkpoint['valid_accuracies']
arch_genotypes = checkpoint['arch_genotypes']
discrepancies = checkpoint['discrepancies']
max_bytes = checkpoint['max_bytes']
logger.log(
"=> loading checkpoint of the last-info '{:}' start with {:}-th epoch."
.format(resume_path, start_epoch))
else:
logger.log(
"=> do not find the last-info file : {:} or resume : {:}".format(
last_info, args.resume))
start_epoch, valid_accuracies, arch_genotypes, discrepancies, max_bytes = 0, {
'best': -1
}, {}, {}, {}
# main procedure
train_func, valid_func = get_procedures(args.procedure)
total_epoch = optim_config.epochs + optim_config.warmup
start_time, epoch_time = time.time(), AverageMeter()
for epoch in range(start_epoch, total_epoch):
scheduler.update(epoch, 0.0)
search_model.set_tau(args.gumbel_tau_max, args.gumbel_tau_min,
epoch * 1.0 / total_epoch)
need_time = 'Time Left: {:}'.format(
convert_secs2time(epoch_time.avg * (total_epoch - epoch), True))
epoch_str = 'epoch={:03d}/{:03d}'.format(epoch, total_epoch)
LRs = scheduler.get_lr()
find_best = False
logger.log(
'\n***{:s}*** start {:s} {:s}, LR=[{:.6f} ~ {:.6f}], scheduler={:}, tau={:}, FLOP={:.2f}'
.format(time_string(), epoch_str, need_time, min(LRs), max(LRs),
scheduler, search_model.tau, MAX_FLOP))
# train for one epoch
train_base_loss, train_arch_loss, train_acc1, train_acc5 = train_func(search_loader, network, criterion, scheduler, base_optimizer, arch_optimizer, optim_config, \
{'epoch-str' : epoch_str, 'FLOP-exp': MAX_FLOP * args.FLOP_ratio,
'FLOP-weight': args.FLOP_weight, 'FLOP-tolerant': MAX_FLOP * args.FLOP_tolerant}, args.print_freq, logger)
# log the results
logger.log(
'***{:s}*** TRAIN [{:}] base-loss = {:.6f}, arch-loss = {:.6f}, accuracy-1 = {:.2f}, accuracy-5 = {:.2f}'
.format(time_string(), epoch_str, train_base_loss, train_arch_loss,
train_acc1, train_acc5))
cur_FLOP, genotype = search_model.get_flop('genotype',
model_config._asdict(),
None)
arch_genotypes[epoch] = genotype
arch_genotypes['last'] = genotype
logger.log('[{:}] genotype : {:}'.format(epoch_str, genotype))
# save the configuration
configure2str(
genotype,
str(
logger.path('log') /
'seed-{:}-temp.config'.format(args.rand_seed)))
arch_info, discrepancy = search_model.get_arch_info()
logger.log(arch_info)
discrepancies[epoch] = discrepancy
logger.log(
'[{:}] FLOP : {:.2f} MB, ratio : {:.4f}, Expected-ratio : {:.4f}, Discrepancy : {:.3f}'
.format(epoch_str, cur_FLOP, cur_FLOP / MAX_FLOP, args.FLOP_ratio,
np.mean(discrepancy)))
#if cur_FLOP/MAX_FLOP > args.FLOP_ratio:
# init_flop_weight = init_flop_weight * args.FLOP_decay
#else:
# init_flop_weight = init_flop_weight / args.FLOP_decay
# evaluate the performance
if (epoch % args.eval_frequency == 0) or (epoch + 1 == total_epoch):
logger.log('-' * 150)
valid_loss, valid_acc1, valid_acc5 = valid_func(
search_valid_loader, network, criterion, epoch_str,
args.print_freq_eval, logger)
valid_accuracies[epoch] = valid_acc1
logger.log(
'***{:s}*** VALID [{:}] loss = {:.6f}, accuracy@1 = {:.2f}, accuracy@5 = {:.2f} | Best-Valid-Acc@1={:.2f}, Error@1={:.2f}'
.format(time_string(), epoch_str, valid_loss, valid_acc1,
valid_acc5, valid_accuracies['best'],
100 - valid_accuracies['best']))
if valid_acc1 > valid_accuracies['best']:
valid_accuracies['best'] = valid_acc1
arch_genotypes['best'] = genotype
find_best = True
logger.log(
'Currently, the best validation accuracy found at {:03d}-epoch :: acc@1={:.2f}, acc@5={:.2f}, error@1={:.2f}, error@5={:.2f}, save into {:}.'
.format(epoch, valid_acc1, valid_acc5, 100 - valid_acc1,
100 - valid_acc5, model_best_path))
# log the GPU memory usage
#num_bytes = torch.cuda.max_memory_allocated( next(network.parameters()).device ) * 1.0
num_bytes = torch.cuda.max_memory_cached(
next(network.parameters()).device) * 1.0
logger.log(
'[GPU-Memory-Usage on {:} is {:} bytes, {:.2f} KB, {:.2f} MB, {:.2f} GB.]'
.format(
next(network.parameters()).device, int(num_bytes),
num_bytes / 1e3, num_bytes / 1e6, num_bytes / 1e9))
max_bytes[epoch] = num_bytes
# save checkpoint
save_path = save_checkpoint(
{
'epoch': epoch,
'args': deepcopy(args),
'max_bytes': deepcopy(max_bytes),
'valid_accuracies': deepcopy(valid_accuracies),
'model-config': model_config._asdict(),
'optim-config': optim_config._asdict(),
'search_model': search_model.state_dict(),
'scheduler': scheduler.state_dict(),
'base_optimizer': base_optimizer.state_dict(),
'arch_optimizer': arch_optimizer.state_dict(),
'arch_genotypes': arch_genotypes,
'discrepancies': discrepancies,
}, model_base_path, logger)
if find_best: copy_checkpoint(model_base_path, model_best_path, logger)
last_info = save_checkpoint(
{
'epoch': epoch,
'args': deepcopy(args),
'last_checkpoint': save_path,
}, logger.path('info'), logger)
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log('')
logger.log('-' * 100)
last_config_path = logger.path('log') / 'seed-{:}-last.config'.format(
args.rand_seed)
configure2str(arch_genotypes['last'], str(last_config_path))
logger.log('save the last config int {:} :\n{:}'.format(
last_config_path, arch_genotypes['last']))
best_arch, valid_acc = arch_genotypes['best'], valid_accuracies['best']
for key, config in arch_genotypes.items():
if key == 'last': continue
FLOP_ratio = config['estimated_FLOP'] / MAX_FLOP
if abs(FLOP_ratio - args.FLOP_ratio) <= args.FLOP_tolerant:
if valid_acc <= valid_accuracies[key]:
best_arch, valid_acc = config, valid_accuracies[key]
print('Best-Arch : {:}\nRatio={:}, Valid-ACC={:}'.format(
best_arch, best_arch['estimated_FLOP'] / MAX_FLOP, valid_acc))
best_config_path = logger.path('log') / 'seed-{:}-best.config'.format(
args.rand_seed)
configure2str(best_arch, str(best_config_path))
logger.log('save the last config int {:} :\n{:}'.format(
best_config_path, best_arch))
logger.log('\n' + '-' * 200)
logger.log(
'Finish training/validation in {:} with Max-GPU-Memory of {:.2f} GB, and save final checkpoint into {:}'
.format(convert_secs2time(epoch_time.sum, True),
max(v for k, v in max_bytes.items()) / 1e9,
logger.path('info')))
logger.close()
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 = 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, config.test_batch_size), xargs.workers)
logger.log(
'||||||| {:10s} ||||||| Search-Loader-Num={:}, Valid-Loader-Num={:}, batch size={:}'
.format(xargs.dataset, len(search_loader), len(valid_loader),
config.batch_size))
logger.log('||||||| {:10s} ||||||| Config={:}'.format(
xargs.dataset, config))
search_space = get_search_spaces('cell', xargs.search_space_name)
model_config = dict2config(
{
'name': 'RANDOM',
'C': xargs.channel,
'N': xargs.num_cells,
'max_nodes': xargs.max_nodes,
'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)
w_optimizer, w_scheduler, criterion = get_optim_scheduler(
search_model.parameters(), config)
logger.log('w-optimizer : {:}'.format(w_optimizer))
logger.log('w-scheduler : {:}'.format(w_scheduler))
logger.log('criterion : {:}'.format(criterion))
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()
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'])
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}, {}
# start training
start_time, search_time, epoch_time, total_epoch = time.time(
), AverageMeter(), AverageMeter(), config.epochs + config.warmup
for epoch in range(start_epoch, total_epoch):
w_scheduler.update(epoch, 0.0)
need_time = 'Time Left: {:}'.format(
convert_secs2time(epoch_time.val * (total_epoch - epoch), True))
epoch_str = '{:03d}-{:03d}'.format(epoch, total_epoch)
logger.log('\n[Search the {:}-th epoch] {:}, LR={:}'.format(
epoch_str, need_time, min(w_scheduler.get_lr())))
# selected_arch = search_find_best(valid_loader, network, criterion, xargs.select_num)
search_w_loss, search_w_top1, search_w_top5 = search_func(
search_loader, network, criterion, w_scheduler, w_optimizer,
epoch_str, xargs.print_freq, 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))
valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
valid_loader, network, criterion)
logger.log(
'[{:}] evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'
.format(epoch_str, valid_a_loss, valid_a_top1, valid_a_top5))
cur_arch, cur_valid_acc = search_find_best(valid_loader, network,
xargs.select_num)
logger.log('[{:}] find-the-best : {:}, accuracy@1={:.2f}%'.format(
epoch_str, cur_arch, cur_valid_acc))
genotypes[epoch] = cur_arch
# check the best accuracy
valid_accuracies[epoch] = valid_a_top1
if valid_a_top1 > valid_accuracies['best']:
valid_accuracies['best'] = valid_a_top1
find_best = True
else:
find_best = False
# save checkpoint
save_path = save_checkpoint(
{
'epoch': epoch + 1,
'args': deepcopy(xargs),
'search_model': search_model.state_dict(),
'w_optimizer': w_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)
if api is not None:
logger.log('{:}'.format(api.query_by_arch(genotypes[epoch])))
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log('\n' + '-' * 200)
logger.log('Pre-searching costs {:.1f} s'.format(search_time.sum))
start_time = time.time()
best_arch, best_acc = search_find_best(valid_loader, network,
xargs.select_num)
search_time.update(time.time() - start_time)
logger.log(
'RANDOM-NAS finds the best one : {:} with accuracy={:.2f}%, with {:.1f} s.'
.format(best_arch, best_acc, search_time.sum))
if api is not None: logger.log('{:}'.format(api.query_by_arch(best_arch)))
logger.close()
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)
# class_num = 4
# xshape = (1,3,88,88)
# print(xshape)
if xargs.dataset == 'cifar10' or xargs.dataset == 'cifar100':
split_Fpath = '/home/city/Projects/NAS-Projects/configs/nas-benchmark/cifar-split.txt'
cifar_split = load_config(split_Fpath, None, None)
train_split, valid_split = cifar_split.train, cifar_split.valid
logger.log('Load split file from {:}'.format(split_Fpath))
elif xargs.dataset.startswith('ImageNet16'):
split_Fpath = 'configs/nas-benchmark/{:}-split.txt'.format(xargs.dataset)
imagenet16_split = load_config(split_Fpath, None, None)
train_split, valid_split = imagenet16_split.train, imagenet16_split.valid
logger.log('Load split file from {:}'.format(split_Fpath))
else:
raise ValueError('invalid dataset : {:}'.format(xargs.dataset))
config_path = '/home/city/Projects/NAS-Projects/configs/nas-benchmark/algos/DARTS.config'
config = load_config(config_path, {'class_num': class_num, 'xshape': xshape}, logger)
print('config')
print(config)
# To split data
train_data_v2 = deepcopy(train_data)
train_data_v2.transform = valid_data.transform
valid_data = train_data_v2
search_data = SearchDataset(xargs.dataset, train_data, train_split, valid_split)
# data loader
search_loader = torch.utils.data.DataLoader(search_data, batch_size=config.batch_size, shuffle=True , num_workers=xargs.workers, pin_memory=True)
valid_loader = torch.utils.data.DataLoader(valid_data, batch_size=config.batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(valid_split), num_workers=xargs.workers, pin_memory=True)
logger.log('||||||| {:10s} ||||||| Search-Loader-Num={:}, Valid-Loader-Num={:}, batch size={:}'.format(xargs.dataset, len(search_loader), len(valid_loader), config.batch_size))
logger.log('||||||| {:10s} ||||||| Config={:}'.format(xargs.dataset, config))
# train_transform = transforms.Compose([
# transforms.RandomHorizontalFlip(),
# transforms.ToTensor()
# # transforms.Normalize(mean=[128, 128, 128], std=[50, 50, 50])
# ])
# val_transform = transforms.Compose([
# # transforms.RandomHorizontalFlip(),
# transforms.ToTensor()
# # transforms.Normalize(mean=[128, 128, 128], std=[50, 50, 50])
# ])
#
# train_data = datasets.ImageFolder(root='/home/city/Projects/build_assessment/data/train',
# transform=train_transform)
# valid_data = datasets.ImageFolder(root='/home/city/Projects/build_assessment/data/val',
# transform=val_transform)
# print(len(train_data))
# print('2333333333333333333333333333333')
# train_split = []
# valid_split = []
#
# for i in range(len(train_data)):
# if i%2==0:
# train_split.append(i)
# else:
# valid_split.append(i)
# search_data = SearchDataset('builds', train_data, train_split, valid_split)
#
# search_loader = torch.utils.data.DataLoader(search_data,
# batch_size=32, shuffle=True,
# num_workers=4, pin_memory=True
# )
# valid_loader = torch.utils.data.DataLoader(valid_data,
# batch_size=32, shuffle=True,
# num_workers=2, pin_memory=True
# )
search_space = get_search_spaces('cell', xargs.search_space_name)
model_config = dict2config({'name': 'DARTS-V2', 'C': xargs.channel, 'N': xargs.num_cells,
'max_nodes': xargs.max_nodes, 'num_classes': class_num,
'space' : search_space}, None)
search_model = get_cell_based_tiny_net(model_config)
logger.log('search-model :\n{:}'.format(search_model))
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))
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()
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}, {}
# start training
start_time, search_time, epoch_time, total_epoch = time.time(), AverageMeter(), AverageMeter(), config.epochs + config.warmup
for epoch in range(start_epoch, total_epoch):
w_scheduler.update(epoch, 0.0)
need_time = 'Time Left: {:}'.format( convert_secs2time(epoch_time.val * (total_epoch-epoch), True) )
epoch_str = '{:03d}-{:03d}'.format(epoch, total_epoch)
min_LR = min(w_scheduler.get_lr())
logger.log('\n[Search the {:}-th epoch] {:}, LR={:}'.format(epoch_str, need_time, min_LR))
search_w_loss, search_w_top1, search_w_top5 = search_func(search_loader, network, criterion, w_scheduler, w_optimizer, a_optimizer, epoch_str, xargs.print_freq, 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))
valid_a_loss , valid_a_top1 , valid_a_top5 = valid_func(valid_loader, network, criterion)
logger.log('[{:}] evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'.format(epoch_str, valid_a_loss, valid_a_top1, valid_a_top5))
# 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()
op_list, _ = genotypes['best'].tolist(remove_str=None)
find_best = True
best_arch_nums = op_list2str(op_list)
torch.save(search_model,'/home/city/disk/log/builds-darts/darts2_%04d_%s_%s_%.2f.pth' %(epoch,time_string_short(),best_arch_nums, valid_a_top1))
print('/home/city/disk/log/builds-darts/darts2_%04d_%s_%s_%.2f.pth' %(epoch,time_string_short(),best_arch_nums, valid_a_top1))
else: find_best = False
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('arch-parameters :\n{:}'.format( nn.functional.softmax(search_model.arch_parameters, dim=-1).cpu() ))
logger.log('arch :\n{:}'.format(nn.functional.softmax(search_model.arch_parameters, dim=-1).cpu().argmax(dim=-1)))
if api is not None: logger.log('{:}'.format(api.query_by_arch( genotypes[epoch] )))
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log('\n' + '-'*100)
# check the performance from the architecture dataset
logger.log('DARTS-V2 : 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()
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, test_data, xshape, class_num = get_datasets(xargs.dataset, xargs.data_path, -1)
if xargs.dataset == 'cifar10' or xargs.dataset == 'cifar100':
split_Fpath = 'configs/nas-benchmark/cifar-split.txt'
cifar_split = load_config(split_Fpath, None, None)
train_split, valid_split = cifar_split.train, cifar_split.valid
logger.log('Load split file from {:}'.format(split_Fpath))
elif xargs.dataset.startswith('ImageNet16'):
split_Fpath = 'configs/nas-benchmark/{:}-split.txt'.format(xargs.dataset)
imagenet16_split = load_config(split_Fpath, None, None)
train_split, valid_split = imagenet16_split.train, imagenet16_split.valid
logger.log('Load split file from {:}'.format(split_Fpath))
else:
raise ValueError('invalid dataset : {:}'.format(xargs.dataset))
logger.log('use config from : {:}'.format(xargs.config_path))
config = load_config(xargs.config_path, {'class_num': class_num, 'xshape': xshape}, logger)
logger.log('config: {:}'.format(config))
# To split data
train_data_v2 = deepcopy(train_data)
train_data_v2.transform = test_data.transform
valid_data = train_data_v2
# data loader
train_loader = torch.utils.data.DataLoader(train_data, batch_size=config.batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(train_split), num_workers=xargs.workers, pin_memory=True)
valid_loader = torch.utils.data.DataLoader(valid_data, batch_size=config.batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(valid_split), num_workers=xargs.workers, pin_memory=True)
logger.log('||||||| {:10s} ||||||| Train-Loader-Num={:}, Valid-Loader-Num={:}, batch size={:}'.format(xargs.dataset, len(train_loader), len(valid_loader), config.batch_size))
logger.log('||||||| {:10s} ||||||| Config={:}'.format(xargs.dataset, config))
search_space = get_search_spaces('cell', xargs.search_space_name)
model_config = dict2config({'name': 'ENAS', 'C': xargs.channel, 'N': xargs.num_cells,
'max_nodes': xargs.max_nodes, 'num_classes': class_num,
'space' : search_space}, None)
shared_cnn = get_cell_based_tiny_net(model_config)
controller = shared_cnn.create_controller()
w_optimizer, w_scheduler, criterion = get_optim_scheduler(shared_cnn.parameters(), config)
a_optimizer = torch.optim.Adam(controller.parameters(), lr=config.controller_lr, betas=config.controller_betas, eps=config.controller_eps)
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(shared_cnn, xshape)
#logger.log('{:}'.format(shared_cnn))
#logger.log('FLOP = {:.2f} M, Params = {:.2f} MB'.format(flop, param))
shared_cnn, controller, criterion = torch.nn.DataParallel(shared_cnn).cuda(), controller.cuda(), criterion.cuda()
last_info, model_base_path, model_best_path = logger.path('info'), logger.path('model'), logger.path('best')
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']
baseline = checkpoint['baseline']
valid_accuracies = checkpoint['valid_accuracies']
shared_cnn.load_state_dict( checkpoint['shared_cnn'] )
controller.load_state_dict( checkpoint['controller'] )
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, baseline = 0, {'best': -1}, {}, None
# start training
start_time, epoch_time, total_epoch = time.time(), AverageMeter(), config.epochs + config.warmup
for epoch in range(start_epoch, total_epoch):
w_scheduler.update(epoch, 0.0)
need_time = 'Time Left: {:}'.format( convert_secs2time(epoch_time.val * (total_epoch-epoch), True) )
epoch_str = '{:03d}-{:03d}'.format(epoch, total_epoch)
logger.log('\n[Search the {:}-th epoch] {:}, LR={:}, baseline={:}'.format(epoch_str, need_time, min(w_scheduler.get_lr()), baseline))
cnn_loss, cnn_top1, cnn_top5 = train_shared_cnn(train_loader, shared_cnn, controller, criterion, w_scheduler, w_optimizer, epoch_str, xargs.print_freq, logger)
logger.log('[{:}] shared-cnn : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'.format(epoch_str, cnn_loss, cnn_top1, cnn_top5))
ctl_loss, ctl_acc, ctl_baseline, ctl_reward, baseline \
= train_controller(valid_loader, shared_cnn, controller, criterion, a_optimizer, \
dict2config({'baseline': baseline,
'ctl_train_steps': xargs.controller_train_steps, 'ctl_num_aggre': xargs.controller_num_aggregate,
'ctl_entropy_w': xargs.controller_entropy_weight,
'ctl_bl_dec' : xargs.controller_bl_dec}, None), \
epoch_str, xargs.print_freq, logger)
logger.log('[{:}] controller : loss={:.2f}, accuracy={:.2f}%, baseline={:.2f}, reward={:.2f}, current-baseline={:.4f}'.format(epoch_str, ctl_loss, ctl_acc, ctl_baseline, ctl_reward, baseline))
best_arch, _ = get_best_arch(controller, shared_cnn, valid_loader)
shared_cnn.module.update_arch(best_arch)
_, best_valid_acc, _ = valid_func(valid_loader, shared_cnn, criterion)
genotypes[epoch] = best_arch
# check the best accuracy
valid_accuracies[epoch] = best_valid_acc
if best_valid_acc > valid_accuracies['best']:
valid_accuracies['best'] = best_valid_acc
genotypes['best'] = best_arch
find_best = True
else: find_best = False
logger.log('<<<--->>> The {:}-th epoch : {:}'.format(epoch_str, genotypes[epoch]))
# save checkpoint
save_path = save_checkpoint({'epoch' : epoch + 1,
'args' : deepcopy(xargs),
'baseline' : baseline,
'shared_cnn' : shared_cnn.state_dict(),
'controller' : controller.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, best_valid_acc))
copy_checkpoint(model_base_path, model_best_path, logger)
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log('\n' + '-'*100)
logger.log('During searching, the best architecture is {:}'.format(genotypes['best']))
logger.log('Its accuracy is {:.2f}%'.format(valid_accuracies['best']))
logger.log('Randomly select {:} architectures and select the best.'.format(xargs.controller_num_samples))
final_arch, _ = get_best_arch(controller, shared_cnn, valid_loader, xargs.controller_num_samples)
shared_cnn.module.update_arch(final_arch)
final_loss, final_top1, final_top5 = valid_func(valid_loader, shared_cnn, criterion)
logger.log('The Selected Final Architecture : {:}'.format(final_arch))
logger.log('Loss={:.3f}, Accuracy@1={:.2f}%, Accuracy@5={:.2f}%'.format(final_loss, final_top1, final_top5))
# check the performance from the architecture dataset
#if xargs.arch_nas_dataset is None or not os.path.isfile(xargs.arch_nas_dataset):
# logger.log('Can not find the architecture dataset : {:}.'.format(xargs.arch_nas_dataset))
#else:
# nas_bench = NASBenchmarkAPI(xargs.arch_nas_dataset)
# geno = genotypes[total_epoch-1]
# logger.log('The last model is {:}'.format(geno))
# info = nas_bench.query_by_arch( geno )
# if info is None: logger.log('Did not find this architecture : {:}.'.format(geno))
# else : logger.log('{:}'.format(info))
# logger.log('-'*100)
# geno = genotypes['best']
# logger.log('The best model is {:}'.format(geno))
# info = nas_bench.query_by_arch( geno )
# if info is None: logger.log('Did not find this architecture : {:}.'.format(geno))
# else : logger.log('{:}'.format(info))
logger.close()
def main(args):
assert torch.cuda.is_available(), "CUDA is not available."
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
torch.set_num_threads(args.workers)
prepare_seed(args.rand_seed)
logger = prepare_logger(args)
train_data, valid_data, xshape, class_num = get_datasets(
args.dataset, args.data_path, args.cutout_length
)
train_loader = torch.utils.data.DataLoader(
train_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
)
valid_loader = torch.utils.data.DataLoader(
valid_data,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
)
# get configures
model_config = load_config(args.model_config, {"class_num": class_num}, logger)
optim_config = load_config(args.optim_config, {"class_num": class_num}, logger)
if args.model_source == "normal":
base_model = obtain_model(model_config)
elif args.model_source == "nas":
base_model = obtain_nas_infer_model(model_config, args.extra_model_path)
elif args.model_source == "autodl-searched":
base_model = obtain_model(model_config, args.extra_model_path)
else:
raise ValueError("invalid model-source : {:}".format(args.model_source))
flop, param = get_model_infos(base_model, xshape)
logger.log("model ====>>>>:\n{:}".format(base_model))
logger.log("model information : {:}".format(base_model.get_message()))
logger.log("-" * 50)
logger.log(
"Params={:.2f} MB, FLOPs={:.2f} M ... = {:.2f} G".format(
param, flop, flop / 1e3
)
)
logger.log("-" * 50)
logger.log("train_data : {:}".format(train_data))
logger.log("valid_data : {:}".format(valid_data))
optimizer, scheduler, criterion = get_optim_scheduler(
base_model.parameters(), optim_config
)
logger.log("optimizer : {:}".format(optimizer))
logger.log("scheduler : {:}".format(scheduler))
logger.log("criterion : {:}".format(criterion))
last_info, model_base_path, model_best_path = (
logger.path("info"),
logger.path("model"),
logger.path("best"),
)
network, criterion = torch.nn.DataParallel(base_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_infox = torch.load(last_info)
start_epoch = last_infox["epoch"] + 1
last_checkpoint_path = last_infox["last_checkpoint"]
if not last_checkpoint_path.exists():
logger.log(
"Does not find {:}, try another path".format(last_checkpoint_path)
)
last_checkpoint_path = (
last_info.parent
/ last_checkpoint_path.parent.name
/ last_checkpoint_path.name
)
checkpoint = torch.load(last_checkpoint_path)
base_model.load_state_dict(checkpoint["base-model"])
scheduler.load_state_dict(checkpoint["scheduler"])
optimizer.load_state_dict(checkpoint["optimizer"])
valid_accuracies = checkpoint["valid_accuracies"]
max_bytes = checkpoint["max_bytes"]
logger.log(
"=> loading checkpoint of the last-info '{:}' start with {:}-th epoch.".format(
last_info, start_epoch
)
)
elif args.resume is not None:
assert Path(args.resume).exists(), "Can not find the resume file : {:}".format(
args.resume
)
checkpoint = torch.load(args.resume)
start_epoch = checkpoint["epoch"] + 1
base_model.load_state_dict(checkpoint["base-model"])
scheduler.load_state_dict(checkpoint["scheduler"])
optimizer.load_state_dict(checkpoint["optimizer"])
valid_accuracies = checkpoint["valid_accuracies"]
max_bytes = checkpoint["max_bytes"]
logger.log(
"=> loading checkpoint from '{:}' start with {:}-th epoch.".format(
args.resume, start_epoch
)
)
elif args.init_model is not None:
assert Path(
args.init_model
).exists(), "Can not find the initialization file : {:}".format(args.init_model)
checkpoint = torch.load(args.init_model)
base_model.load_state_dict(checkpoint["base-model"])
start_epoch, valid_accuracies, max_bytes = 0, {"best": -1}, {}
logger.log("=> initialize the model from {:}".format(args.init_model))
else:
logger.log("=> do not find the last-info file : {:}".format(last_info))
start_epoch, valid_accuracies, max_bytes = 0, {"best": -1}, {}
train_func, valid_func = get_procedures(args.procedure)
total_epoch = optim_config.epochs + optim_config.warmup
# Main Training and Evaluation Loop
start_time = time.time()
epoch_time = AverageMeter()
for epoch in range(start_epoch, total_epoch):
scheduler.update(epoch, 0.0)
need_time = "Time Left: {:}".format(
convert_secs2time(epoch_time.avg * (total_epoch - epoch), True)
)
epoch_str = "epoch={:03d}/{:03d}".format(epoch, total_epoch)
LRs = scheduler.get_lr()
find_best = False
# set-up drop-out ratio
if hasattr(base_model, "update_drop_path"):
base_model.update_drop_path(
model_config.drop_path_prob * epoch / total_epoch
)
logger.log(
"\n***{:s}*** start {:s} {:s}, LR=[{:.6f} ~ {:.6f}], scheduler={:}".format(
time_string(), epoch_str, need_time, min(LRs), max(LRs), scheduler
)
)
# train for one epoch
train_loss, train_acc1, train_acc5 = train_func(
train_loader,
network,
criterion,
scheduler,
optimizer,
optim_config,
epoch_str,
args.print_freq,
logger,
)
# log the results
logger.log(
"***{:s}*** TRAIN [{:}] loss = {:.6f}, accuracy-1 = {:.2f}, accuracy-5 = {:.2f}".format(
time_string(), epoch_str, train_loss, train_acc1, train_acc5
)
)
# evaluate the performance
if (epoch % args.eval_frequency == 0) or (epoch + 1 == total_epoch):
logger.log("-" * 150)
valid_loss, valid_acc1, valid_acc5 = valid_func(
valid_loader,
network,
criterion,
optim_config,
epoch_str,
args.print_freq_eval,
logger,
)
valid_accuracies[epoch] = valid_acc1
logger.log(
"***{:s}*** VALID [{:}] loss = {:.6f}, accuracy@1 = {:.2f}, accuracy@5 = {:.2f} | Best-Valid-Acc@1={:.2f}, Error@1={:.2f}".format(
time_string(),
epoch_str,
valid_loss,
valid_acc1,
valid_acc5,
valid_accuracies["best"],
100 - valid_accuracies["best"],
)
)
if valid_acc1 > valid_accuracies["best"]:
valid_accuracies["best"] = valid_acc1
find_best = True
logger.log(
"Currently, the best validation accuracy found at {:03d}-epoch :: acc@1={:.2f}, acc@5={:.2f}, error@1={:.2f}, error@5={:.2f}, save into {:}.".format(
epoch,
valid_acc1,
valid_acc5,
100 - valid_acc1,
100 - valid_acc5,
model_best_path,
)
)
num_bytes = (
torch.cuda.max_memory_cached(next(network.parameters()).device) * 1.0
)
logger.log(
"[GPU-Memory-Usage on {:} is {:} bytes, {:.2f} KB, {:.2f} MB, {:.2f} GB.]".format(
next(network.parameters()).device,
int(num_bytes),
num_bytes / 1e3,
num_bytes / 1e6,
num_bytes / 1e9,
)
)
max_bytes[epoch] = num_bytes
if epoch % 10 == 0:
torch.cuda.empty_cache()
# save checkpoint
save_path = save_checkpoint(
{
"epoch": epoch,
"args": deepcopy(args),
"max_bytes": deepcopy(max_bytes),
"FLOP": flop,
"PARAM": param,
"valid_accuracies": deepcopy(valid_accuracies),
"model-config": model_config._asdict(),
"optim-config": optim_config._asdict(),
"base-model": base_model.state_dict(),
"scheduler": scheduler.state_dict(),
"optimizer": optimizer.state_dict(),
},
model_base_path,
logger,
)
if find_best:
copy_checkpoint(model_base_path, model_best_path, logger)
last_info = save_checkpoint(
{
"epoch": epoch,
"args": deepcopy(args),
"last_checkpoint": save_path,
},
logger.path("info"),
logger,
)
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log("\n" + "-" * 200)
logger.log(
"Finish training/validation in {:} with Max-GPU-Memory of {:.2f} MB, and save final checkpoint into {:}".format(
convert_secs2time(epoch_time.sum, True),
max(v for k, v in max_bytes.items()) / 1e6,
logger.path("info"),
)
)
logger.log("-" * 200 + "\n")
logger.close()
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)
if xargs.dataset == 'cifar10' or xargs.dataset == 'cifar100':
split_Fpath = 'configs/nas-benchmark/cifar-split.txt'
cifar_split = load_config(split_Fpath, None, None)
train_split, valid_split = cifar_split.train, cifar_split.valid
logger.log('Load split file from {:}'.format(split_Fpath))
#elif xargs.dataset.startswith('ImageNet16'):
# # all_indexes = list(range(len(train_data))) ; random.seed(111) ; random.shuffle(all_indexes)
# # train_split, valid_split = sorted(all_indexes[: len(train_data)//2]), sorted(all_indexes[len(train_data)//2 :])
# # imagenet16_split = dict2config({'train': train_split, 'valid': valid_split}, None)
# # _ = configure2str(imagenet16_split, 'temp.txt')
# split_Fpath = 'configs/nas-benchmark/{:}-split.txt'.format(xargs.dataset)
# imagenet16_split = load_config(split_Fpath, None, None)
# train_split, valid_split = imagenet16_split.train, imagenet16_split.valid
# logger.log('Load split file from {:}'.format(split_Fpath))
else:
raise ValueError('invalid dataset : {:}'.format(xargs.dataset))
config = load_config(xargs.config_path, {
'class_num': class_num,
'xshape': xshape
}, logger)
logger.log('config : {:}'.format(config))
# To split data
train_data_v2 = deepcopy(train_data)
train_data_v2.transform = valid_data.transform
valid_data = train_data_v2
search_data = SearchDataset(xargs.dataset, train_data, train_split,
valid_split)
# data loader
search_loader = torch.utils.data.DataLoader(search_data,
batch_size=config.batch_size,
shuffle=True,
num_workers=xargs.workers,
pin_memory=True)
valid_loader = torch.utils.data.DataLoader(
valid_data,
batch_size=config.batch_size,
sampler=torch.utils.data.sampler.SubsetRandomSampler(valid_split),
num_workers=xargs.workers,
pin_memory=True)
logger.log(
'||||||| {:10s} ||||||| Search-Loader-Num={:}, Valid-Loader-Num={:}, batch size={:}'
.format(xargs.dataset, len(search_loader), len(valid_loader),
config.batch_size))
logger.log('||||||| {:10s} ||||||| Config={:}'.format(
xargs.dataset, config))
search_space = get_search_spaces('cell', xargs.search_space_name)
model_config = dict2config(
{
'name': 'RANDOM',
'C': xargs.channel,
'N': xargs.num_cells,
'max_nodes': xargs.max_nodes,
'num_classes': class_num,
'space': search_space
}, None)
search_model = get_cell_based_tiny_net(model_config)
w_optimizer, w_scheduler, criterion = get_optim_scheduler(
search_model.parameters(), config)
logger.log('w-optimizer : {:}'.format(w_optimizer))
logger.log('w-scheduler : {:}'.format(w_scheduler))
logger.log('criterion : {:}'.format(criterion))
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()
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'])
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'])
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 = 0, {'best': -1}
# start training
start_time, epoch_time, total_epoch = time.time(), AverageMeter(
), config.epochs + config.warmup
for epoch in range(start_epoch, total_epoch):
w_scheduler.update(epoch, 0.0)
need_time = 'Time Left: {:}'.format(
convert_secs2time(epoch_time.val * (total_epoch - epoch), True))
epoch_str = '{:03d}-{:03d}'.format(epoch, total_epoch)
logger.log('\n[Search the {:}-th epoch] {:}, LR={:}'.format(
epoch_str, need_time, min(w_scheduler.get_lr())))
search_w_loss, search_w_top1, search_w_top5 = search_func(
search_loader, network, criterion, w_scheduler, w_optimizer,
epoch_str, xargs.print_freq, logger)
logger.log(
'[{:}] searching : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'
.format(epoch_str, search_w_loss, search_w_top1, search_w_top5))
valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
valid_loader, network, criterion)
logger.log(
'[{:}] evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'
.format(epoch_str, valid_a_loss, valid_a_top1, valid_a_top5))
# check the best accuracy
valid_accuracies[epoch] = valid_a_top1
if valid_a_top1 > valid_accuracies['best']:
valid_accuracies['best'] = valid_a_top1
find_best = True
else:
find_best = False
# save checkpoint
save_path = save_checkpoint(
{
'epoch': epoch + 1,
'args': deepcopy(xargs),
'search_model': search_model.state_dict(),
'w_optimizer': w_optimizer.state_dict(),
'w_scheduler': w_scheduler.state_dict(),
'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)
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log('\n' + '-' * 200)
best_arch, best_acc = None, -1
for iarch in range(xargs.select_num):
arch = search_model.random_genotype(True)
valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
valid_loader, network, criterion)
logger.log(
'final evaluation [{:02d}/{:02d}] : {:} : accuracy={:.2f}%, loss={:.3f}'
.format(iarch, xargs.select_num, arch, valid_a_top1, valid_a_loss))
if best_arch is None or best_acc < valid_a_top1:
best_arch, best_acc = arch, valid_a_top1
logger.log('Find the best one : {:} with accuracy={:.2f}%'.format(
best_arch, best_acc))
logger.log('\n' + '-' * 100)
"""
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, test_data, xshape, class_num = get_datasets(
xargs.dataset, xargs.data_path, -1)
logger.log("use config from : {:}".format(xargs.config_path))
config = load_config(xargs.config_path, {
"class_num": class_num,
"xshape": xshape
}, logger)
_, train_loader, valid_loader = get_nas_search_loaders(
train_data,
test_data,
xargs.dataset,
"configs/nas-benchmark/",
config.batch_size,
xargs.workers,
)
# since ENAS will train the controller on valid-loader, we need to use train transformation for valid-loader
valid_loader.dataset.transform = deepcopy(train_loader.dataset.transform)
if hasattr(valid_loader.dataset, "transforms"):
valid_loader.dataset.transforms = deepcopy(
train_loader.dataset.transforms)
# data loader
logger.log(
"||||||| {:10s} ||||||| Train-Loader-Num={:}, Valid-Loader-Num={:}, batch size={:}"
.format(xargs.dataset, len(train_loader), len(valid_loader),
config.batch_size))
logger.log("||||||| {:10s} ||||||| Config={:}".format(
xargs.dataset, config))
search_space = get_search_spaces("cell", xargs.search_space_name)
model_config = dict2config(
{
"name": "ENAS",
"C": xargs.channel,
"N": xargs.num_cells,
"max_nodes": xargs.max_nodes,
"num_classes": class_num,
"space": search_space,
"affine": False,
"track_running_stats": bool(xargs.track_running_stats),
},
None,
)
shared_cnn = get_cell_based_tiny_net(model_config)
controller = shared_cnn.create_controller()
w_optimizer, w_scheduler, criterion = get_optim_scheduler(
shared_cnn.parameters(), config)
a_optimizer = torch.optim.Adam(
controller.parameters(),
lr=config.controller_lr,
betas=config.controller_betas,
eps=config.controller_eps,
)
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(shared_cnn, xshape)
# logger.log('{:}'.format(shared_cnn))
# logger.log('FLOP = {:.2f} M, Params = {:.2f} MB'.format(flop, param))
logger.log("search-space : {:}".format(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))
shared_cnn, controller, criterion = (
torch.nn.DataParallel(shared_cnn).cuda(),
controller.cuda(),
criterion.cuda(),
)
last_info, model_base_path, model_best_path = (
logger.path("info"),
logger.path("model"),
logger.path("best"),
)
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"]
baseline = checkpoint["baseline"]
valid_accuracies = checkpoint["valid_accuracies"]
shared_cnn.load_state_dict(checkpoint["shared_cnn"])
controller.load_state_dict(checkpoint["controller"])
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, baseline = 0, {
"best": -1
}, {}, None
# start training
start_time, search_time, epoch_time, total_epoch = (
time.time(),
AverageMeter(),
AverageMeter(),
config.epochs + config.warmup,
)
for epoch in range(start_epoch, total_epoch):
w_scheduler.update(epoch, 0.0)
need_time = "Time Left: {:}".format(
convert_secs2time(epoch_time.val * (total_epoch - epoch), True))
epoch_str = "{:03d}-{:03d}".format(epoch, total_epoch)
logger.log(
"\n[Search the {:}-th epoch] {:}, LR={:}, baseline={:}".format(
epoch_str, need_time, min(w_scheduler.get_lr()), baseline))
cnn_loss, cnn_top1, cnn_top5 = train_shared_cnn(
train_loader,
shared_cnn,
controller,
criterion,
w_scheduler,
w_optimizer,
epoch_str,
xargs.print_freq,
logger,
)
logger.log(
"[{:}] shared-cnn : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%"
.format(epoch_str, cnn_loss, cnn_top1, cnn_top5))
ctl_loss, ctl_acc, ctl_baseline, ctl_reward, baseline = train_controller(
valid_loader,
shared_cnn,
controller,
criterion,
a_optimizer,
dict2config(
{
"baseline": baseline,
"ctl_train_steps": xargs.controller_train_steps,
"ctl_num_aggre": xargs.controller_num_aggregate,
"ctl_entropy_w": xargs.controller_entropy_weight,
"ctl_bl_dec": xargs.controller_bl_dec,
},
None,
),
epoch_str,
xargs.print_freq,
logger,
)
search_time.update(time.time() - start_time)
logger.log(
"[{:}] controller : loss={:.2f}, accuracy={:.2f}%, baseline={:.2f}, reward={:.2f}, current-baseline={:.4f}, time-cost={:.1f} s"
.format(
epoch_str,
ctl_loss,
ctl_acc,
ctl_baseline,
ctl_reward,
baseline,
search_time.sum,
))
best_arch, _ = get_best_arch(controller, shared_cnn, valid_loader)
shared_cnn.module.update_arch(best_arch)
_, best_valid_acc, _ = valid_func(valid_loader, shared_cnn, criterion)
genotypes[epoch] = best_arch
# check the best accuracy
valid_accuracies[epoch] = best_valid_acc
if best_valid_acc > valid_accuracies["best"]:
valid_accuracies["best"] = best_valid_acc
genotypes["best"] = best_arch
find_best = True
else:
find_best = False
logger.log("<<<--->>> The {:}-th epoch : {:}".format(
epoch_str, genotypes[epoch]))
# save checkpoint
save_path = save_checkpoint(
{
"epoch": epoch + 1,
"args": deepcopy(xargs),
"baseline": baseline,
"shared_cnn": shared_cnn.state_dict(),
"controller": controller.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, best_valid_acc))
copy_checkpoint(model_base_path, model_best_path, logger)
if api is not None:
logger.log("{:}".format(api.query_by_arch(genotypes[epoch],
"200")))
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log("\n" + "-" * 100)
logger.log("During searching, the best architecture is {:}".format(
genotypes["best"]))
logger.log("Its accuracy is {:.2f}%".format(valid_accuracies["best"]))
logger.log("Randomly select {:} architectures and select the best.".format(
xargs.controller_num_samples))
start_time = time.time()
final_arch, _ = get_best_arch(controller, shared_cnn, valid_loader,
xargs.controller_num_samples)
search_time.update(time.time() - start_time)
shared_cnn.module.update_arch(final_arch)
final_loss, final_top1, final_top5 = valid_func(valid_loader, shared_cnn,
criterion)
logger.log("The Selected Final Architecture : {:}".format(final_arch))
logger.log("Loss={:.3f}, Accuracy@1={:.2f}%, Accuracy@5={:.2f}%".format(
final_loss, final_top1, final_top5))
logger.log(
"ENAS : run {:} epochs, cost {:.1f} s, last-geno is {:}.".format(
total_epoch, search_time.sum, final_arch))
if api is not None:
logger.log("{:}".format(api.query_by_arch(final_arch)))
logger.close()
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 = 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, config.test_batch_size),
xargs.workers,
)
logger.log(
"||||||| {:10s} ||||||| Search-Loader-Num={:}, Valid-Loader-Num={:}, batch size={:}"
.format(xargs.dataset, len(search_loader), len(valid_loader),
config.batch_size))
logger.log("||||||| {:10s} ||||||| Config={:}".format(
xargs.dataset, config))
search_space = get_search_spaces("cell", xargs.search_space_name)
model_config = dict2config(
{
"name": "RANDOM",
"C": xargs.channel,
"N": xargs.num_cells,
"max_nodes": xargs.max_nodes,
"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)
w_optimizer, w_scheduler, criterion = get_optim_scheduler(
search_model.parameters(), config)
logger.log("w-optimizer : {:}".format(w_optimizer))
logger.log("w-scheduler : {:}".format(w_scheduler))
logger.log("criterion : {:}".format(criterion))
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()
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"])
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}, {}
# start training
start_time, search_time, epoch_time, total_epoch = (
time.time(),
AverageMeter(),
AverageMeter(),
config.epochs + config.warmup,
)
for epoch in range(start_epoch, total_epoch):
w_scheduler.update(epoch, 0.0)
need_time = "Time Left: {:}".format(
convert_secs2time(epoch_time.val * (total_epoch - epoch), True))
epoch_str = "{:03d}-{:03d}".format(epoch, total_epoch)
logger.log("\n[Search the {:}-th epoch] {:}, LR={:}".format(
epoch_str, need_time, min(w_scheduler.get_lr())))
# selected_arch = search_find_best(valid_loader, network, criterion, xargs.select_num)
search_w_loss, search_w_top1, search_w_top5 = search_func(
search_loader,
network,
criterion,
w_scheduler,
w_optimizer,
epoch_str,
xargs.print_freq,
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))
valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
valid_loader, network, criterion)
logger.log(
"[{:}] evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%"
.format(epoch_str, valid_a_loss, valid_a_top1, valid_a_top5))
cur_arch, cur_valid_acc = search_find_best(valid_loader, network,
xargs.select_num)
logger.log("[{:}] find-the-best : {:}, accuracy@1={:.2f}%".format(
epoch_str, cur_arch, cur_valid_acc))
genotypes[epoch] = cur_arch
# check the best accuracy
valid_accuracies[epoch] = valid_a_top1
if valid_a_top1 > valid_accuracies["best"]:
valid_accuracies["best"] = valid_a_top1
find_best = True
else:
find_best = False
# save checkpoint
save_path = save_checkpoint(
{
"epoch": epoch + 1,
"args": deepcopy(xargs),
"search_model": search_model.state_dict(),
"w_optimizer": w_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)
if api is not None:
logger.log("{:}".format(api.query_by_arch(genotypes[epoch],
"200")))
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log("\n" + "-" * 200)
logger.log("Pre-searching costs {:.1f} s".format(search_time.sum))
start_time = time.time()
best_arch, best_acc = search_find_best(valid_loader, network,
xargs.select_num)
search_time.update(time.time() - start_time)
logger.log(
"RANDOM-NAS finds the best one : {:} with accuracy={:.2f}%, with {:.1f} s."
.format(best_arch, best_acc, search_time.sum))
if api is not None:
logger.log("{:}".format(api.query_by_arch(best_arch, "200")))
logger.close()
def main(args):
assert torch.cuda.is_available(), "CUDA is not available."
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
torch.set_num_threads(args.workers)
prepare_seed(args.rand_seed)
logger = prepare_logger(args)
# prepare dataset
train_data, valid_data, xshape, class_num = get_datasets(
args.dataset, args.data_path, args.cutout_length)
# train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True , num_workers=args.workers, pin_memory=True)
valid_loader = torch.utils.data.DataLoader(
valid_data,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
)
split_file_path = Path(args.split_path)
assert split_file_path.exists(), "{:} does not exist".format(
split_file_path)
split_info = torch.load(split_file_path)
train_split, valid_split = split_info["train"], split_info["valid"]
assert (len(set(train_split).intersection(set(valid_split))) == 0
), "There should be 0 element that belongs to both train and valid"
assert len(train_split) + len(valid_split) == len(
train_data), "{:} + {:} vs {:}".format(len(train_split),
len(valid_split),
len(train_data))
search_dataset = SearchDataset(args.dataset, train_data, train_split,
valid_split)
search_train_loader = torch.utils.data.DataLoader(
train_data,
batch_size=args.batch_size,
sampler=torch.utils.data.sampler.SubsetRandomSampler(train_split),
pin_memory=True,
num_workers=args.workers,
)
search_valid_loader = torch.utils.data.DataLoader(
train_data,
batch_size=args.batch_size,
sampler=torch.utils.data.sampler.SubsetRandomSampler(valid_split),
pin_memory=True,
num_workers=args.workers,
)
search_loader = torch.utils.data.DataLoader(
search_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
sampler=None,
)
# get configures
model_config = load_config(
args.model_config,
{
"class_num": class_num,
"search_mode": args.search_shape
},
logger,
)
# obtain the model
search_model = obtain_search_model(model_config)
MAX_FLOP, param = get_model_infos(search_model, xshape)
optim_config = load_config(args.optim_config, {
"class_num": class_num,
"FLOP": MAX_FLOP
}, logger)
logger.log("Model Information : {:}".format(search_model.get_message()))
logger.log("MAX_FLOP = {:} M".format(MAX_FLOP))
logger.log("Params = {:} M".format(param))
logger.log("train_data : {:}".format(train_data))
logger.log("search-data: {:}".format(search_dataset))
logger.log("search_train_loader : {:} samples".format(len(train_split)))
logger.log("search_valid_loader : {:} samples".format(len(valid_split)))
base_optimizer, scheduler, criterion = get_optim_scheduler(
search_model.base_parameters(), optim_config)
arch_optimizer = torch.optim.Adam(
search_model.arch_parameters(),
lr=optim_config.arch_LR,
betas=(0.5, 0.999),
weight_decay=optim_config.arch_decay,
)
logger.log("base-optimizer : {:}".format(base_optimizer))
logger.log("arch-optimizer : {:}".format(arch_optimizer))
logger.log("scheduler : {:}".format(scheduler))
logger.log("criterion : {:}".format(criterion))
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()
# load checkpoint
if last_info.exists() or (args.resume is not None and osp.isfile(
args.resume)): # automatically resume from previous checkpoint
if args.resume is not None and osp.isfile(args.resume):
resume_path = Path(args.resume)
elif last_info.exists():
resume_path = last_info
else:
raise ValueError("Something is wrong.")
logger.log("=> loading checkpoint of the last-info '{:}' start".format(
resume_path))
checkpoint = torch.load(resume_path)
if "last_checkpoint" in checkpoint:
last_checkpoint_path = checkpoint["last_checkpoint"]
if not last_checkpoint_path.exists():
logger.log("Does not find {:}, try another path".format(
last_checkpoint_path))
last_checkpoint_path = (resume_path.parent /
last_checkpoint_path.parent.name /
last_checkpoint_path.name)
assert (last_checkpoint_path.exists()
), "can not find the checkpoint from {:}".format(
last_checkpoint_path)
checkpoint = torch.load(last_checkpoint_path)
start_epoch = checkpoint["epoch"] + 1
search_model.load_state_dict(checkpoint["search_model"])
scheduler.load_state_dict(checkpoint["scheduler"])
base_optimizer.load_state_dict(checkpoint["base_optimizer"])
arch_optimizer.load_state_dict(checkpoint["arch_optimizer"])
valid_accuracies = checkpoint["valid_accuracies"]
arch_genotypes = checkpoint["arch_genotypes"]
discrepancies = checkpoint["discrepancies"]
logger.log(
"=> loading checkpoint of the last-info '{:}' start with {:}-th epoch."
.format(resume_path, start_epoch))
else:
logger.log(
"=> do not find the last-info file : {:} or resume : {:}".format(
last_info, args.resume))
start_epoch, valid_accuracies, arch_genotypes, discrepancies = (
0,
{
"best": -1
},
{},
{},
)
# main procedure
train_func, valid_func = get_procedures(args.procedure)
total_epoch = optim_config.epochs + optim_config.warmup
start_time, epoch_time = time.time(), AverageMeter()
for epoch in range(start_epoch, total_epoch):
scheduler.update(epoch, 0.0)
search_model.set_tau(args.gumbel_tau_max, args.gumbel_tau_min,
epoch * 1.0 / total_epoch)
need_time = "Time Left: {:}".format(
convert_secs2time(epoch_time.avg * (total_epoch - epoch), True))
epoch_str = "epoch={:03d}/{:03d}".format(epoch, total_epoch)
LRs = scheduler.get_lr()
find_best = False
logger.log(
"\n***{:s}*** start {:s} {:s}, LR=[{:.6f} ~ {:.6f}], scheduler={:}, tau={:}, FLOP={:.2f}"
.format(
time_string(),
epoch_str,
need_time,
min(LRs),
max(LRs),
scheduler,
search_model.tau,
MAX_FLOP,
))
# train for one epoch
train_base_loss, train_arch_loss, train_acc1, train_acc5 = train_func(
search_loader,
network,
criterion,
scheduler,
base_optimizer,
arch_optimizer,
optim_config,
{
"epoch-str": epoch_str,
"FLOP-exp": MAX_FLOP * args.FLOP_ratio,
"FLOP-weight": args.FLOP_weight,
"FLOP-tolerant": MAX_FLOP * args.FLOP_tolerant,
},
args.print_freq,
logger,
)
# log the results
logger.log(
"***{:s}*** TRAIN [{:}] base-loss = {:.6f}, arch-loss = {:.6f}, accuracy-1 = {:.2f}, accuracy-5 = {:.2f}"
.format(
time_string(),
epoch_str,
train_base_loss,
train_arch_loss,
train_acc1,
train_acc5,
))
cur_FLOP, genotype = search_model.get_flop("genotype",
model_config._asdict(),
None)
arch_genotypes[epoch] = genotype
arch_genotypes["last"] = genotype
logger.log("[{:}] genotype : {:}".format(epoch_str, genotype))
arch_info, discrepancy = search_model.get_arch_info()
logger.log(arch_info)
discrepancies[epoch] = discrepancy
logger.log(
"[{:}] FLOP : {:.2f} MB, ratio : {:.4f}, Expected-ratio : {:.4f}, Discrepancy : {:.3f}"
.format(
epoch_str,
cur_FLOP,
cur_FLOP / MAX_FLOP,
args.FLOP_ratio,
np.mean(discrepancy),
))
# if cur_FLOP/MAX_FLOP > args.FLOP_ratio:
# init_flop_weight = init_flop_weight * args.FLOP_decay
# else:
# init_flop_weight = init_flop_weight / args.FLOP_decay
# evaluate the performance
if (epoch % args.eval_frequency == 0) or (epoch + 1 == total_epoch):
logger.log("-" * 150)
valid_loss, valid_acc1, valid_acc5 = valid_func(
search_valid_loader,
network,
criterion,
epoch_str,
args.print_freq_eval,
logger,
)
valid_accuracies[epoch] = valid_acc1
logger.log(
"***{:s}*** VALID [{:}] loss = {:.6f}, accuracy@1 = {:.2f}, accuracy@5 = {:.2f} | Best-Valid-Acc@1={:.2f}, Error@1={:.2f}"
.format(
time_string(),
epoch_str,
valid_loss,
valid_acc1,
valid_acc5,
valid_accuracies["best"],
100 - valid_accuracies["best"],
))
if valid_acc1 > valid_accuracies["best"]:
valid_accuracies["best"] = valid_acc1
arch_genotypes["best"] = genotype
find_best = True
logger.log(
"Currently, the best validation accuracy found at {:03d}-epoch :: acc@1={:.2f}, acc@5={:.2f}, error@1={:.2f}, error@5={:.2f}, save into {:}."
.format(
epoch,
valid_acc1,
valid_acc5,
100 - valid_acc1,
100 - valid_acc5,
model_best_path,
))
# save checkpoint
save_path = save_checkpoint(
{
"epoch": epoch,
"args": deepcopy(args),
"valid_accuracies": deepcopy(valid_accuracies),
"model-config": model_config._asdict(),
"optim-config": optim_config._asdict(),
"search_model": search_model.state_dict(),
"scheduler": scheduler.state_dict(),
"base_optimizer": base_optimizer.state_dict(),
"arch_optimizer": arch_optimizer.state_dict(),
"arch_genotypes": arch_genotypes,
"discrepancies": discrepancies,
},
model_base_path,
logger,
)
if find_best:
copy_checkpoint(model_base_path, model_best_path, logger)
last_info = save_checkpoint(
{
"epoch": epoch,
"args": deepcopy(args),
"last_checkpoint": save_path,
},
logger.path("info"),
logger,
)
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log("")
logger.log("-" * 100)
last_config_path = logger.path("log") / "seed-{:}-last.config".format(
args.rand_seed)
configure2str(arch_genotypes["last"], str(last_config_path))
logger.log("save the last config int {:} :\n{:}".format(
last_config_path, arch_genotypes["last"]))
best_arch, valid_acc = arch_genotypes["best"], valid_accuracies["best"]
for key, config in arch_genotypes.items():
if key == "last":
continue
FLOP_ratio = config["estimated_FLOP"] / MAX_FLOP
if abs(FLOP_ratio - args.FLOP_ratio) <= args.FLOP_tolerant:
if valid_acc < valid_accuracies[key]:
best_arch, valid_acc = config, valid_accuracies[key]
print("Best-Arch : {:}\nRatio={:}, Valid-ACC={:}".format(
best_arch, best_arch["estimated_FLOP"] / MAX_FLOP, valid_acc))
best_config_path = logger.path("log") / "seed-{:}-best.config".format(
args.rand_seed)
configure2str(best_arch, str(best_config_path))
logger.log("save the last config int {:} :\n{:}".format(
best_config_path, best_arch))
logger.log("\n" + "-" * 200)
logger.log(
"Finish training/validation in {:}, and save final checkpoint into {:}"
.format(convert_secs2time(epoch_time.sum, True), logger.path("info")))
logger.close()
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': 'GDAS',
'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': 'GDAS',
'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:
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(
'GDAS : 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()
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:
model_config = dict2config(
{
'name': 'GDAS',
'C': xargs.channel,
'N': xargs.num_cells,
'max_nodes': xargs.max_nodes,
'num_classes': class_num,
'space': search_space,
'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('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()
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
for epoch in range(start_epoch, total_epoch):
w_scheduler.update(epoch, 0.0)
need_time = 'Time Left: {:}'.format(
convert_secs2time(epoch_time.val * (total_epoch - epoch), 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())))
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)
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))
# 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
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:
logger.log('{:}'.format(api.query_by_arch(genotypes[epoch],
'200')))
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log('\n' + '-' * 100)
# check the performance from the architecture dataset
logger.log(
'GDAS : 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], '200')))
logger.close()
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, test_data, xshape, class_num = get_datasets(
xargs.dataset, xargs.data_path, -1)
print(xshape)
xshape = (1, 3, 88, 88)
class_num = 4
if xargs.dataset == 'cifar10' or xargs.dataset == 'cifar100':
split_Fpath = '/home/city/Projects/NAS-Projects/configs/nas-benchmark/cifar-split.txt'
cifar_split = load_config(split_Fpath, None, None)
train_split, valid_split = cifar_split.train, cifar_split.valid
logger.log('Load split file from {:}'.format(split_Fpath))
# elif xargs.dataset.startswith('ImageNet16'):
# split_Fpath = '/home/city/Projects/NAS-Projects/configs/nas-benchmark/{:}-split.txt'.format(xargs.dataset)
# imagenet16_split = load_config(split_Fpath, None, None)
# train_split, valid_split = imagenet16_split.train, imagenet16_split.valid
# logger.log('Load split file from {:}'.format(split_Fpath))
# else:
# raise ValueError('invalid dataset : {:}'.format(xargs.dataset))
logger.log('use config from : {:}'.format(xargs.config_path))
config = load_config(xargs.config_path, {
'class_num': class_num,
'xshape': xshape
}, logger)
logger.log('config: {:}'.format(config))
# To split data
train_data_v2 = deepcopy(train_data)
train_data_v2.transform = test_data.transform
valid_data = train_data_v2
train_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[128, 128, 128], std=[50, 50, 50])
])
val_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[128, 128, 128], std=[50, 50, 50])
])
train_data = datasets.ImageFolder(
root='/home/city/Projects/build_assessment/data/train',
transform=train_transform)
valid_data = datasets.ImageFolder(
root='/home/city/Projects/build_assessment/data/val',
transform=val_transform)
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=32,
shuffle=True,
num_workers=4)
valid_loader = torch.utils.data.DataLoader(valid_data,
batch_size=32,
shuffle=True,
num_workers=2)
# data loader
# train_loader = torch.utils.data.DataLoader(train_data, batch_size=config.batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(train_split), num_workers=xargs.workers, pin_memory=True)
# valid_loader = torch.utils.data.DataLoader(valid_data, batch_size=config.batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(valid_split), num_workers=xargs.workers, pin_memory=True)
logger.log(
'||||||| {:10s} ||||||| Train-Loader-Num={:}, Valid-Loader-Num={:}, batch size={:}'
.format(xargs.dataset, len(train_loader), len(valid_loader),
config.batch_size))
logger.log('||||||| {:10s} ||||||| Config={:}'.format(
xargs.dataset, config))
search_space = get_search_spaces('cell', xargs.search_space_name)
print('search space', search_space)
model_config = dict2config(
{
'name': 'ENAS',
'C': xargs.channel,
'N': xargs.num_cells,
'max_nodes': xargs.max_nodes,
'num_classes': class_num,
'space': search_space
}, None)
shared_cnn = get_cell_based_tiny_net(model_config)
# input_data = Variable(torch.rand(16, 3, 224, 224))
# print(shared_cnn)
# writer = SummaryWriter(log_dir='/home/city/disk/tmp/log', comment='resnet18')
# with writer:
# writer.add_graph(shared_cnn, (input_data,))
# exit()
print(len(shared_cnn.edge2index), len(shared_cnn.op_names),
'lend edge2index,len op names')
print(shared_cnn.edge2index)
print(shared_cnn.op_names)
controller = shared_cnn.create_controller()
print(controller)
w_optimizer, w_scheduler, criterion = get_optim_scheduler(
shared_cnn.parameters(), config)
a_optimizer = torch.optim.Adam(controller.parameters(),
lr=config.controller_lr,
betas=config.controller_betas,
eps=config.controller_eps)
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(shared_cnn, xshape)
#logger.log('{:}'.format(shared_cnn))
#logger.log('FLOP = {:.2f} M, Params = {:.2f} MB'.format(flop, param))
logger.log('search-space : {:}'.format(search_space))
# if xargs.arch_nas_dataset is None:
# # api = None
# # else:
# # api = API(xargs.arch_nas_dataset)
api = None
logger.log('{:} create API = {:} done'.format(time_string(), api))
shared_cnn, controller, criterion = torch.nn.DataParallel(
shared_cnn).cuda(), controller.cuda(), criterion.cuda()
#tw.model_stats(shared_cnn, [13, 3, 224, 224])
last_info, model_base_path, model_best_path = logger.path(
'info'), logger.path('model'), logger.path('best')
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']
baseline = checkpoint['baseline']
valid_accuracies = checkpoint['valid_accuracies']
shared_cnn.load_state_dict(checkpoint['shared_cnn'])
controller.load_state_dict(checkpoint['controller'])
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, baseline = 0, {
'best': -1
}, {}, None
pre_train_shared_cnn(train_loader, valid_loader, shared_cnn, controller,
criterion, w_scheduler, w_optimizer)
w_optimizer, w_scheduler, criterion = get_optim_scheduler(
shared_cnn.parameters(), config)
# start training
start_time, search_time, epoch_time, total_epoch = time.time(
), AverageMeter(), AverageMeter(), config.epochs + config.warmup
for epoch in range(start_epoch, total_epoch):
w_scheduler.update(epoch, 0.0)
need_time = 'Time Left: {:}'.format(
convert_secs2time(epoch_time.val * (total_epoch - epoch), True))
epoch_str = '{:03d}-{:03d}'.format(epoch, total_epoch)
logger.log(
'\n[Search the {:}-th epoch] {:}, LR={:}, baseline={:}'.format(
epoch_str, need_time, min(w_scheduler.get_lr()), baseline))
cnn_loss, cnn_top1, cnn_top5 = train_shared_cnn(
train_loader, shared_cnn, controller, criterion, w_scheduler,
w_optimizer, epoch_str, xargs.print_freq, logger)
logger.log(
'[{:}] shared-cnn : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'
.format(epoch_str, cnn_loss, cnn_top1, cnn_top5))
ctl_loss, ctl_acc, ctl_baseline, ctl_reward, baseline \
= train_controller(valid_loader, shared_cnn, controller, criterion, a_optimizer, \
dict2config({'baseline': baseline,
'ctl_train_steps': xargs.controller_train_steps, 'ctl_num_aggre': xargs.controller_num_aggregate,
'ctl_entropy_w': xargs.controller_entropy_weight,
'ctl_bl_dec' : xargs.controller_bl_dec}, None), \
epoch_str, xargs.print_freq, logger)
search_time.update(time.time() - start_time)
logger.log(
'[{:}] controller : loss={:.2f}, accuracy={:.2f}%, baseline={:.2f}, reward={:.2f}, current-baseline={:.4f}, time-cost={:.1f} s'
.format(epoch_str, ctl_loss, ctl_acc, ctl_baseline, ctl_reward,
baseline, search_time.sum))
best_arch, _ = get_best_arch(controller, shared_cnn, valid_loader)
shared_cnn.module.update_arch(best_arch)
_, best_valid_acc, _ = valid_func(valid_loader, shared_cnn, criterion)
op_list, _ = best_arch.tolist(remove_str=None)
best_arch_nums = op_list2str(op_list)
genotypes[epoch] = best_arch
# check the best accuracy
valid_accuracies[epoch] = best_valid_acc
if best_valid_acc > valid_accuracies['best']:
valid_accuracies['best'] = best_valid_acc
genotypes['best'] = best_arch
find_best = True
torch.save(
shared_cnn,
'/home/city/disk/log/builds-enas/share_cnn_%s_%s_%.2f.pth' %
(time_string_short(), best_arch_nums, best_valid_acc))
torch.save(
controller,
'/home/city/disk/log/builds-enas/controller_%s_%s_%.2f.pth' %
(time_string_short(), best_arch_nums, best_valid_acc))
print('/home/city/disk/log/builds-enas/share_cnn_%s_%s_%.2f.pth' %
(time_string_short(), best_arch_nums, best_valid_acc))
else:
find_best = False
print(best_valid_acc, valid_accuracies['best'])
logger.log('<<<--->>> The {:}-th epoch : {:}'.format(
epoch_str, genotypes[epoch]))
# save checkpoint
save_path = save_checkpoint(
{
'epoch': epoch + 1,
'args': deepcopy(xargs),
'baseline': baseline,
'shared_cnn': shared_cnn.state_dict(),
'controller': controller.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, best_valid_acc))
copy_checkpoint(model_base_path, model_best_path, logger)
if api is not None:
logger.log('{:}'.format(api.query_by_arch(genotypes[epoch])))
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
# with open("/home/city/disk/tmp/foo.txt") as file:
# in_x = torch.rand(13, 3, 128, 128)
# # in_x =
# g = make_dot(shared_cnn(in_x), params=dict(list(shared_cnn.named_parameters()) + [('x', in_x)]))
# g.render('/home/city/disk/tmp/espnet_model_%04d' % epoch, view=False)
# # g = None
logger.log('\n' + '-' * 100)
logger.log('During searching, the best architecture is {:}'.format(
genotypes['best']))
logger.log('Its accuracy is {:.2f}%'.format(valid_accuracies['best']))
logger.log('Randomly select {:} architectures and select the best.'.format(
xargs.controller_num_samples))
start_time = time.time()
final_arch, _ = get_best_arch(controller, shared_cnn, valid_loader,
xargs.controller_num_samples)
search_time.update(time.time() - start_time)
shared_cnn.module.update_arch(final_arch)
final_loss, final_top1, final_top5 = valid_func(valid_loader, shared_cnn,
criterion)
logger.log('The Selected Final Architecture : {:}'.format(final_arch))
logger.log('Loss={:.3f}, Accuracy@1={:.2f}%, Accuracy@5={:.2f}%'.format(
final_loss, final_top1, final_top5))
logger.log(
'ENAS : run {:} epochs, cost {:.1f} s, last-geno is {:}.'.format(
total_epoch, search_time.sum, final_arch))
if api is not None: logger.log('{:}'.format(api.query_by_arch(final_arch)))
logger.close()