def online_evaluate(
env,
meta_model,
base_model,
criterion,
metric,
args,
logger,
save=False,
easy_adapt=False,
):
logger.log("Online evaluate: {:}".format(env))
metric.reset()
loss_meter = AverageMeter()
w_containers = dict()
for idx, (future_time, (future_x, future_y)) in enumerate(env):
with torch.no_grad():
meta_model.eval()
base_model.eval()
future_time_embed = meta_model.gen_time_embed(
future_time.to(args.device).view(-1))
[future_container] = meta_model.gen_model(future_time_embed)
if save:
w_containers[idx] = future_container.no_grad_clone()
future_x, future_y = future_x.to(args.device), future_y.to(
args.device)
future_y_hat = base_model.forward_with_container(
future_x, future_container)
future_loss = criterion(future_y_hat, future_y)
loss_meter.update(future_loss.item())
# accumulate the metric scores
score = metric(future_y_hat, future_y)
if easy_adapt:
meta_model.easy_adapt(future_time.item(), future_time_embed)
refine, post_refine_loss = False, -1
else:
refine, post_refine_loss = meta_model.adapt(
base_model,
criterion,
future_time.item(),
future_x,
future_y,
args.refine_lr,
args.refine_epochs,
{
"param": future_time_embed,
"loss": future_loss.item()
},
)
logger.log(
"[ONLINE] [{:03d}/{:03d}] loss={:.4f}, score={:.4f}".format(
idx, len(env), future_loss.item(), score) +
", post-loss={:.4f}".format(post_refine_loss if refine else -1))
meta_model.clear_fixed()
meta_model.clear_learnt()
return w_containers, loss_meter.avg, metric.get_info()["score"]
def search_valid(xloader, network, criterion, extra_info, print_freq, logger):
data_time, batch_time, losses, top1, top5 = (
AverageMeter(),
AverageMeter(),
AverageMeter(),
AverageMeter(),
AverageMeter(),
)
network.eval()
network.apply(change_key("search_mode", "search"))
end = time.time()
# logger.log('Starting evaluating {:}'.format(epoch_info))
with torch.no_grad():
for i, (inputs, targets) in enumerate(xloader):
# measure data loading time
data_time.update(time.time() - end)
# calculate prediction and loss
targets = targets.cuda(non_blocking=True)
logits, expected_flop = network(inputs)
loss = criterion(logits, targets)
# record
prec1, prec5 = obtain_accuracy(logits.data,
targets.data,
topk=(1, 5))
losses.update(loss.item(), inputs.size(0))
top1.update(prec1.item(), inputs.size(0))
top5.update(prec5.item(), inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % print_freq == 0 or (i + 1) == len(xloader):
Sstr = ("**VALID** " + time_string() +
" [{:}][{:03d}/{:03d}]".format(extra_info, i,
len(xloader)))
Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format(
batch_time=batch_time, data_time=data_time)
Lstr = "Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})".format(
loss=losses, top1=top1, top5=top5)
Istr = "Size={:}".format(list(inputs.size()))
logger.log(Sstr + " " + Tstr + " " + Lstr + " " + Istr)
logger.log(
" **VALID** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Loss:{loss:.3f}"
.format(
top1=top1,
top5=top5,
error1=100 - top1.avg,
error5=100 - top5.avg,
loss=losses.avg,
))
return losses.avg, top1.avg, top5.avg
def procedure(xloader, network, criterion, scheduler, optimizer, mode: str):
losses, top1, top5 = AverageMeter(), AverageMeter(), AverageMeter()
if mode == "train":
network.train()
elif mode == "valid":
network.eval()
else:
raise ValueError("The mode is not right : {:}".format(mode))
device = torch.cuda.current_device()
data_time, batch_time, end = AverageMeter(), AverageMeter(), time.time()
for i, (inputs, targets) in enumerate(xloader):
if mode == "train":
scheduler.update(None, 1.0 * i / len(xloader))
targets = targets.cuda(device=device, non_blocking=True)
if mode == "train":
optimizer.zero_grad()
# forward
features, logits = network(inputs)
loss = criterion(logits, targets)
# backward
if mode == "train":
loss.backward()
optimizer.step()
# record loss and accuracy
prec1, prec5 = obtain_accuracy(logits.data, targets.data, topk=(1, 5))
losses.update(loss.item(), inputs.size(0))
top1.update(prec1.item(), inputs.size(0))
top5.update(prec5.item(), inputs.size(0))
# count time
batch_time.update(time.time() - end)
end = time.time()
return losses.avg, top1.avg, top5.avg, batch_time.sum
def train_shared_cnn(
xloader,
shared_cnn,
controller,
criterion,
scheduler,
optimizer,
epoch_str,
print_freq,
logger,
):
data_time, batch_time = AverageMeter(), AverageMeter()
losses, top1s, top5s, xend = (
AverageMeter(),
AverageMeter(),
AverageMeter(),
time.time(),
)
shared_cnn.train()
controller.eval()
for step, (inputs, targets) in enumerate(xloader):
scheduler.update(None, 1.0 * step / len(xloader))
targets = targets.cuda(non_blocking=True)
# measure data loading time
data_time.update(time.time() - xend)
with torch.no_grad():
_, _, sampled_arch = controller()
optimizer.zero_grad()
shared_cnn.module.update_arch(sampled_arch)
_, logits = shared_cnn(inputs)
loss = criterion(logits, targets)
loss.backward()
torch.nn.utils.clip_grad_norm_(shared_cnn.parameters(), 5)
optimizer.step()
# record
prec1, prec5 = obtain_accuracy(logits.data, targets.data, topk=(1, 5))
losses.update(loss.item(), inputs.size(0))
top1s.update(prec1.item(), inputs.size(0))
top5s.update(prec5.item(), inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - xend)
xend = time.time()
if step % print_freq == 0 or step + 1 == len(xloader):
Sstr = (
"*Train-Shared-CNN* " + time_string() +
" [{:}][{:03d}/{:03d}]".format(epoch_str, step, len(xloader)))
Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format(
batch_time=batch_time, data_time=data_time)
Wstr = "[Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})]".format(
loss=losses, top1=top1s, top5=top5s)
logger.log(Sstr + " " + Tstr + " " + Wstr)
return losses.avg, top1s.avg, top5s.avg
def procedure(
xloader,
network,
criterion,
optimizer,
metric,
mode: Text,
logger_fn: Callable = None,
):
data_time, batch_time = AverageMeter(), AverageMeter()
if mode.lower() == "train":
network.train()
elif mode.lower() == "valid":
network.eval()
else:
raise ValueError("The mode is not right : {:}".format(mode))
end = time.time()
for i, (inputs, targets) in enumerate(xloader):
# measure data loading time
data_time.update(time.time() - end)
# calculate prediction and loss
if mode == "train":
optimizer.zero_grad()
outputs = network(inputs)
targets = targets.to(get_device(outputs))
if mode == "train":
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
# record
with torch.no_grad():
results = metric(outputs, targets)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
return metric.get_info()
def test_func(
xloader,
network,
criterion,
):
base_losses, base_top1, base_top5 = AverageMeter(), AverageMeter(), AverageMeter()
network.eval()
for step, (base_inputs, base_targets) in enumerate(
xloader
):
base_targets = base_targets.cuda(non_blocking=True)
_, logits = network(base_inputs.cuda())
base_loss = criterion(logits, base_targets)
base_prec1, base_prec5 = obtain_accuracy(
logits.data, base_targets.data, topk=(1, 5)
)
base_losses.update(base_loss.item(), base_inputs.size(0))
base_top1.update(base_prec1.item(), base_inputs.size(0))
base_top5.update(base_prec5.item(), base_inputs.size(0))
return (
base_losses.avg,
base_top1.avg,
base_top5.avg,
)
def search_func(
xloader,
network,
criterion,
scheduler,
w_optimizer,
a_optimizer,
epoch_str,
print_freq,
logger,
):
data_time, batch_time = AverageMeter(), AverageMeter()
base_losses, base_top1, base_top5 = AverageMeter(), AverageMeter(
), AverageMeter()
arch_losses, arch_top1, arch_top5 = AverageMeter(), AverageMeter(
), AverageMeter()
end = time.time()
network.train()
for step, (base_inputs, base_targets, arch_inputs,
arch_targets) in enumerate(xloader):
scheduler.update(None, 1.0 * step / len(xloader))
base_targets = base_targets.cuda(non_blocking=True)
arch_targets = arch_targets.cuda(non_blocking=True)
# measure data loading time
data_time.update(time.time() - end)
# update the weights
sampled_arch = network.module.dync_genotype(True)
network.module.set_cal_mode("dynamic", sampled_arch)
# network.module.set_cal_mode( 'urs' )
network.zero_grad()
_, logits = network(base_inputs)
base_loss = criterion(logits, base_targets)
base_loss.backward()
w_optimizer.step()
# record
base_prec1, base_prec5 = obtain_accuracy(logits.data,
base_targets.data,
topk=(1, 5))
base_losses.update(base_loss.item(), base_inputs.size(0))
base_top1.update(base_prec1.item(), base_inputs.size(0))
base_top5.update(base_prec5.item(), base_inputs.size(0))
# update the architecture-weight
network.module.set_cal_mode("joint")
network.zero_grad()
_, logits = network(arch_inputs)
arch_loss = criterion(logits, arch_targets)
arch_loss.backward()
a_optimizer.step()
# record
arch_prec1, arch_prec5 = obtain_accuracy(logits.data,
arch_targets.data,
topk=(1, 5))
arch_losses.update(arch_loss.item(), arch_inputs.size(0))
arch_top1.update(arch_prec1.item(), arch_inputs.size(0))
arch_top5.update(arch_prec5.item(), arch_inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if step % print_freq == 0 or step + 1 == len(xloader):
Sstr = (
"*SEARCH* " + time_string() +
" [{:}][{:03d}/{:03d}]".format(epoch_str, step, len(xloader)))
Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format(
batch_time=batch_time, data_time=data_time)
Wstr = "Base [Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})]".format(
loss=base_losses, top1=base_top1, top5=base_top5)
Astr = "Arch [Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})]".format(
loss=arch_losses, top1=arch_top1, top5=arch_top5)
logger.log(Sstr + " " + Tstr + " " + Wstr + " " + Astr)
# print (nn.functional.softmax(network.module.arch_parameters, dim=-1))
# print (network.module.arch_parameters)
return (
base_losses.avg,
base_top1.avg,
base_top5.avg,
arch_losses.avg,
arch_top1.avg,
arch_top5.avg,
)
def search_train_v2(
search_loader,
network,
criterion,
scheduler,
base_optimizer,
arch_optimizer,
optim_config,
extra_info,
print_freq,
logger,
):
data_time, batch_time = AverageMeter(), AverageMeter()
base_losses, arch_losses, top1, top5 = (
AverageMeter(),
AverageMeter(),
AverageMeter(),
AverageMeter(),
)
arch_cls_losses, arch_flop_losses = AverageMeter(), AverageMeter()
epoch_str, flop_need, flop_weight, flop_tolerant = (
extra_info["epoch-str"],
extra_info["FLOP-exp"],
extra_info["FLOP-weight"],
extra_info["FLOP-tolerant"],
)
network.train()
logger.log(
"[Search] : {:}, FLOP-Require={:.2f} MB, FLOP-WEIGHT={:.2f}".format(
epoch_str, flop_need, flop_weight
)
)
end = time.time()
network.apply(change_key("search_mode", "search"))
for step, (base_inputs, base_targets, arch_inputs, arch_targets) in enumerate(
search_loader
):
scheduler.update(None, 1.0 * step / len(search_loader))
# calculate prediction and loss
base_targets = base_targets.cuda(non_blocking=True)
arch_targets = arch_targets.cuda(non_blocking=True)
# measure data loading time
data_time.update(time.time() - end)
# update the weights
base_optimizer.zero_grad()
logits, expected_flop = network(base_inputs)
base_loss = criterion(logits, base_targets)
base_loss.backward()
base_optimizer.step()
# record
prec1, prec5 = obtain_accuracy(logits.data, base_targets.data, topk=(1, 5))
base_losses.update(base_loss.item(), base_inputs.size(0))
top1.update(prec1.item(), base_inputs.size(0))
top5.update(prec5.item(), base_inputs.size(0))
# update the architecture
arch_optimizer.zero_grad()
logits, expected_flop = network(arch_inputs)
flop_cur = network.module.get_flop("genotype", None, None)
flop_loss, flop_loss_scale = get_flop_loss(
expected_flop, flop_cur, flop_need, flop_tolerant
)
acls_loss = criterion(logits, arch_targets)
arch_loss = acls_loss + flop_loss * flop_weight
arch_loss.backward()
arch_optimizer.step()
# record
arch_losses.update(arch_loss.item(), arch_inputs.size(0))
arch_flop_losses.update(flop_loss_scale, arch_inputs.size(0))
arch_cls_losses.update(acls_loss.item(), arch_inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if step % print_freq == 0 or (step + 1) == len(search_loader):
Sstr = (
"**TRAIN** "
+ time_string()
+ " [{:}][{:03d}/{:03d}]".format(epoch_str, step, len(search_loader))
)
Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format(
batch_time=batch_time, data_time=data_time
)
Lstr = "Base-Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})".format(
loss=base_losses, top1=top1, top5=top5
)
Vstr = "Acls-loss {aloss.val:.3f} ({aloss.avg:.3f}) FLOP-Loss {floss.val:.3f} ({floss.avg:.3f}) Arch-Loss {loss.val:.3f} ({loss.avg:.3f})".format(
aloss=arch_cls_losses, floss=arch_flop_losses, loss=arch_losses
)
logger.log(Sstr + " " + Tstr + " " + Lstr + " " + Vstr)
# num_bytes = torch.cuda.max_memory_allocated( next(network.parameters()).device ) * 1.0
# logger.log(Sstr + ' ' + Tstr + ' ' + Lstr + ' ' + Vstr + ' GPU={:.2f}MB'.format(num_bytes/1e6))
# Istr = 'Bsz={:} Asz={:}'.format(list(base_inputs.size()), list(arch_inputs.size()))
# logger.log(Sstr + ' ' + Tstr + ' ' + Lstr + ' ' + Vstr + ' ' + Istr)
# print(network.module.get_arch_info())
# print(network.module.width_attentions[0])
# print(network.module.width_attentions[1])
logger.log(
" **TRAIN** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Base-Loss:{baseloss:.3f}, Arch-Loss={archloss:.3f}".format(
top1=top1,
top5=top5,
error1=100 - top1.avg,
error5=100 - top5.avg,
baseloss=base_losses.avg,
archloss=arch_losses.avg,
)
)
return base_losses.avg, arch_losses.avg, top1.avg, top5.avg
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)
if xargs.model_config is None:
model_config = dict2config(
dict(
name="SETN",
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,
dict(
num_classes=class_num,
space=search_space,
affine=False,
track_running_stats=bool(xargs.track_running_stats),
),
None,
)
logger.log("search space : {:}".format(search_space))
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("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))
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))
init_genotype, _ = get_best_arch(valid_loader, network,
xargs.select_num)
start_epoch, valid_accuracies, genotypes = 0, {
"best": -1
}, {
-1: init_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)
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,
)
search_time.update(time.time() - start_time)
logger.log(
"[{:}] search [base] : 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(
"[{:}] search [arch] : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%"
.format(epoch_str, search_a_loss, search_a_top1, search_a_top5))
genotype, temp_accuracy = get_best_arch(valid_loader, network,
xargs.select_num)
network.module.set_cal_mode("dynamic", genotype)
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,
genotype))
# search_model.set_cal_mode('urs')
# valid_a_loss , valid_a_top1 , valid_a_top5 = valid_func(valid_loader, network, criterion)
# logger.log('[{:}] URS---evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'.format(epoch_str, valid_a_loss, valid_a_top1, valid_a_top5))
# search_model.set_cal_mode('joint')
# valid_a_loss , valid_a_top1 , valid_a_top5 = valid_func(valid_loader, network, criterion)
# logger.log('[{:}] JOINT-evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'.format(epoch_str, valid_a_loss, valid_a_top1, valid_a_top5))
# search_model.set_cal_mode('select')
# valid_a_loss , valid_a_top1 , valid_a_top5 = valid_func(valid_loader, network, criterion)
# logger.log('[{:}] Selec-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
genotypes[epoch] = 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,
)
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()
# the final post procedure : count the time
start_time = time.time()
genotype, temp_accuracy = get_best_arch(valid_loader, network,
xargs.select_num)
search_time.update(time.time() - start_time)
network.module.set_cal_mode("dynamic", genotype)
valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
valid_loader, network, criterion)
logger.log(
"Last : the gentotype is : {:}, with the validation accuracy of {:.3f}%."
.format(genotype, valid_a_top1))
logger.log("\n" + "-" * 100)
# check the performance from the architecture dataset
logger.log(
"SETN : run {:} epochs, cost {:.1f} s, last-geno is {:}.".format(
total_epoch, search_time.sum, genotype))
if api is not None:
logger.log("{:}".format(api.query_by_arch(genotype, "200")))
logger.close()
def main(args):
prepare_seed(args.rand_seed)
logger = prepare_logger(args)
train_env = get_synthetic_env(mode="train", version=args.env_version)
valid_env = get_synthetic_env(mode="valid", version=args.env_version)
trainval_env = get_synthetic_env(mode="trainval", version=args.env_version)
test_env = get_synthetic_env(mode="test", version=args.env_version)
all_env = get_synthetic_env(mode=None, version=args.env_version)
logger.log("The training enviornment: {:}".format(train_env))
logger.log("The validation enviornment: {:}".format(valid_env))
logger.log("The trainval enviornment: {:}".format(trainval_env))
logger.log("The total enviornment: {:}".format(all_env))
logger.log("The test enviornment: {:}".format(test_env))
model_kwargs = dict(
config=dict(model_type="norm_mlp"),
input_dim=all_env.meta_info["input_dim"],
output_dim=all_env.meta_info["output_dim"],
hidden_dims=[args.hidden_dim] * 2,
act_cls="relu",
norm_cls="layer_norm_1d",
)
model = get_model(**model_kwargs)
model = model.to(args.device)
if all_env.meta_info["task"] == "regression":
criterion = torch.nn.MSELoss()
metric_cls = MSEMetric
elif all_env.meta_info["task"] == "classification":
criterion = torch.nn.CrossEntropyLoss()
metric_cls = Top1AccMetric
else:
raise ValueError(
"This task ({:}) is not supported.".format(all_env.meta_info["task"])
)
maml = MAML(
model, criterion, args.epochs, args.meta_lr, args.inner_lr, args.inner_step
)
# meta-training
last_success_epoch = 0
per_epoch_time, start_time = AverageMeter(), time.time()
for iepoch in range(args.epochs):
need_time = "Time Left: {:}".format(
convert_secs2time(per_epoch_time.avg * (args.epochs - iepoch), True)
)
head_str = (
"[{:}] [{:04d}/{:04d}] ".format(time_string(), iepoch, args.epochs)
+ need_time
)
maml.zero_grad()
meta_losses = []
for ibatch in range(args.meta_batch):
future_idx = random.randint(0, len(trainval_env) - 1)
future_t, (future_x, future_y) = trainval_env[future_idx]
# -->>
seq_times = trainval_env.get_seq_times(future_idx, args.seq_length)
_, (allxs, allys) = trainval_env.seq_call(seq_times)
allxs, allys = allxs.view(-1, allxs.shape[-1]), allys.view(-1, 1)
if trainval_env.meta_info["task"] == "classification":
allys = allys.view(-1)
historical_x, historical_y = allxs.to(args.device), allys.to(args.device)
future_container = maml.adapt(historical_x, historical_y)
future_x, future_y = future_x.to(args.device), future_y.to(args.device)
future_y_hat = maml.predict(future_x, future_container)
future_loss = maml.criterion(future_y_hat, future_y)
meta_losses.append(future_loss)
meta_loss = torch.stack(meta_losses).mean()
meta_loss.backward()
maml.step()
logger.log(head_str + " meta-loss: {:.4f}".format(meta_loss.item()))
success, best_score = maml.save_best(-meta_loss.item())
if success:
logger.log("Achieve the best with best_score = {:.3f}".format(best_score))
save_checkpoint(maml.state_dict(), logger.path("model"), logger)
last_success_epoch = iepoch
if iepoch - last_success_epoch >= args.early_stop_thresh:
logger.log("Early stop at {:}".format(iepoch))
break
per_epoch_time.update(time.time() - start_time)
start_time = time.time()
# meta-test
maml.load_best()
def finetune(index):
seq_times = test_env.get_seq_times(index, args.seq_length)
_, (allxs, allys) = test_env.seq_call(seq_times)
allxs, allys = allxs.view(-1, allxs.shape[-1]), allys.view(-1, 1)
if test_env.meta_info["task"] == "classification":
allys = allys.view(-1)
historical_x, historical_y = allxs.to(args.device), allys.to(args.device)
future_container = maml.adapt(historical_x, historical_y)
historical_y_hat = maml.predict(historical_x, future_container)
train_metric = metric_cls(True)
# model.analyze_weights()
with torch.no_grad():
train_metric(historical_y_hat, historical_y)
train_results = train_metric.get_info()
return train_results, future_container
metric = metric_cls(True)
per_timestamp_time, start_time = AverageMeter(), time.time()
for idx, (future_time, (future_x, future_y)) in enumerate(test_env):
need_time = "Time Left: {:}".format(
convert_secs2time(per_timestamp_time.avg * (len(test_env) - idx), True)
)
logger.log(
"[{:}]".format(time_string())
+ " [{:04d}/{:04d}]".format(idx, len(test_env))
+ " "
+ need_time
)
# build optimizer
train_results, future_container = finetune(idx)
future_x, future_y = future_x.to(args.device), future_y.to(args.device)
future_y_hat = maml.predict(future_x, future_container)
future_loss = criterion(future_y_hat, future_y)
metric(future_y_hat, future_y)
log_str = (
"[{:}]".format(time_string())
+ " [{:04d}/{:04d}]".format(idx, len(test_env))
+ " train-score: {:.5f}, eval-score: {:.5f}".format(
train_results["score"], metric.get_info()["score"]
)
)
logger.log(log_str)
logger.log("")
per_timestamp_time.update(time.time() - start_time)
start_time = time.time()
logger.log("-" * 200 + "\n")
logger.close()
def train_controller(
xloader,
shared_cnn,
controller,
criterion,
optimizer,
config,
epoch_str,
print_freq,
logger,
):
# config. (containing some necessary arg)
# baseline: The baseline score (i.e. average val_acc) from the previous epoch
data_time, batch_time = AverageMeter(), AverageMeter()
(
GradnormMeter,
LossMeter,
ValAccMeter,
EntropyMeter,
BaselineMeter,
RewardMeter,
xend,
) = (
AverageMeter(),
AverageMeter(),
AverageMeter(),
AverageMeter(),
AverageMeter(),
AverageMeter(),
time.time(),
)
shared_cnn.eval()
controller.train()
controller.zero_grad()
# for step, (inputs, targets) in enumerate(xloader):
loader_iter = iter(xloader)
for step in range(config.ctl_train_steps * config.ctl_num_aggre):
try:
inputs, targets = next(loader_iter)
except:
loader_iter = iter(xloader)
inputs, targets = next(loader_iter)
targets = targets.cuda(non_blocking=True)
# measure data loading time
data_time.update(time.time() - xend)
log_prob, entropy, sampled_arch = controller()
with torch.no_grad():
shared_cnn.module.update_arch(sampled_arch)
_, logits = shared_cnn(inputs)
val_top1, val_top5 = obtain_accuracy(logits.data,
targets.data,
topk=(1, 5))
val_top1 = val_top1.view(-1) / 100
reward = val_top1 + config.ctl_entropy_w * entropy
if config.baseline is None:
baseline = val_top1
else:
baseline = config.baseline - (1 - config.ctl_bl_dec) * (
config.baseline - reward)
loss = -1 * log_prob * (reward - baseline)
# account
RewardMeter.update(reward.item())
BaselineMeter.update(baseline.item())
ValAccMeter.update(val_top1.item() * 100)
LossMeter.update(loss.item())
EntropyMeter.update(entropy.item())
# Average gradient over controller_num_aggregate samples
loss = loss / config.ctl_num_aggre
loss.backward(retain_graph=True)
# measure elapsed time
batch_time.update(time.time() - xend)
xend = time.time()
if (step + 1) % config.ctl_num_aggre == 0:
grad_norm = torch.nn.utils.clip_grad_norm_(controller.parameters(),
5.0)
GradnormMeter.update(grad_norm)
optimizer.step()
controller.zero_grad()
if step % print_freq == 0:
Sstr = ("*Train-Controller* " + time_string() +
" [{:}][{:03d}/{:03d}]".format(
epoch_str, step,
config.ctl_train_steps * config.ctl_num_aggre))
Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format(
batch_time=batch_time, data_time=data_time)
Wstr = "[Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Reward {reward.val:.2f} ({reward.avg:.2f})] Baseline {basel.val:.2f} ({basel.avg:.2f})".format(
loss=LossMeter,
top1=ValAccMeter,
reward=RewardMeter,
basel=BaselineMeter,
)
Estr = "Entropy={:.4f} ({:.4f})".format(EntropyMeter.val,
EntropyMeter.avg)
logger.log(Sstr + " " + Tstr + " " + Wstr + " " + Estr)
return (
LossMeter.avg,
ValAccMeter.avg,
BaselineMeter.avg,
RewardMeter.avg,
baseline.item(),
)
def check_files(save_dir, meta_file, basestr):
meta_infos = torch.load(meta_file, map_location="cpu")
meta_archs = meta_infos["archs"]
meta_num_archs = meta_infos["total"]
assert meta_num_archs == len(
meta_archs), "invalid number of archs : {:} vs {:}".format(
meta_num_archs, len(meta_archs))
sub_model_dirs = sorted(list(save_dir.glob("*-*-{:}".format(basestr))))
print("{:} find {:} directories used to save checkpoints".format(
time_string(), len(sub_model_dirs)))
subdir2archs, num_evaluated_arch = collections.OrderedDict(), 0
num_seeds = defaultdict(lambda: 0)
for index, sub_dir in enumerate(sub_model_dirs):
xcheckpoints = list(sub_dir.glob("arch-*-seed-*.pth"))
# xcheckpoints = list(sub_dir.glob('arch-*-seed-0777.pth')) + list(sub_dir.glob('arch-*-seed-0888.pth')) + list(sub_dir.glob('arch-*-seed-0999.pth'))
arch_indexes = set()
for checkpoint in xcheckpoints:
temp_names = checkpoint.name.split("-")
assert (len(temp_names) == 4 and temp_names[0] == "arch"
and temp_names[2]
== "seed"), "invalid checkpoint name : {:}".format(
checkpoint.name)
arch_indexes.add(temp_names[1])
subdir2archs[sub_dir] = sorted(list(arch_indexes))
num_evaluated_arch += len(arch_indexes)
# count number of seeds for each architecture
for arch_index in arch_indexes:
num_seeds[len(
list(sub_dir.glob(
"arch-{:}-seed-*.pth".format(arch_index))))] += 1
print(
"There are {:5d} architectures that have been evaluated ({:} in total, {:} ckps in total)."
.format(num_evaluated_arch, meta_num_archs,
sum(k * v for k, v in num_seeds.items())))
for key in sorted(list(num_seeds.keys())):
print("There are {:5d} architectures that are evaluated {:} times.".
format(num_seeds[key], key))
dir2ckps, dir2ckp_exists = dict(), dict()
start_time, epoch_time = time.time(), AverageMeter()
for IDX, (sub_dir, arch_indexes) in enumerate(subdir2archs.items()):
if basestr == "C16-N5":
seeds = [777, 888, 999]
elif basestr == "C16-N5-LESS":
seeds = [111, 777]
else:
raise ValueError("Invalid base str : {:}".format(basestr))
numrs = defaultdict(lambda: 0)
all_checkpoints, all_ckp_exists = [], []
for arch_index in arch_indexes:
checkpoints = [
"arch-{:}-seed-{:04d}.pth".format(arch_index, seed)
for seed in seeds
]
ckp_exists = [(sub_dir / x).exists() for x in checkpoints]
arch_index = int(arch_index)
assert (
0 <= arch_index < len(meta_archs)
), "invalid arch-index {:} (not found in meta_archs)".format(
arch_index)
all_checkpoints += checkpoints
all_ckp_exists += ckp_exists
numrs[sum(ckp_exists)] += 1
dir2ckps[str(sub_dir)] = all_checkpoints
dir2ckp_exists[str(sub_dir)] = all_ckp_exists
# measure time
epoch_time.update(time.time() - start_time)
start_time = time.time()
numrstr = ", ".join(
["{:}: {:03d}".format(x, numrs[x]) for x in sorted(numrs.keys())])
print(
"{:} load [{:2d}/{:2d}] [{:03d} archs] [{:04d}->{:04d} ckps] {:} done, need {:}. {:}"
.format(
time_string(),
IDX + 1,
len(subdir2archs),
len(arch_indexes),
len(all_checkpoints),
sum(all_ckp_exists),
sub_dir,
convert_secs2time(
epoch_time.avg * (len(subdir2archs) - IDX - 1), True),
numrstr,
))
def main(
save_dir,
workers,
datasets,
xpaths,
splits,
use_less,
srange,
arch_index,
seeds,
cover_mode,
meta_info,
arch_config,
):
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(workers)
assert (len(srange) == 2 and
0 <= srange[0] <= srange[1]), "invalid srange : {:}".format(srange)
if use_less:
sub_dir = Path(save_dir) / "{:06d}-{:06d}-C{:}-N{:}-LESS".format(
srange[0], srange[1], arch_config["channel"],
arch_config["num_cells"])
else:
sub_dir = Path(save_dir) / "{:06d}-{:06d}-C{:}-N{:}".format(
srange[0], srange[1], arch_config["channel"],
arch_config["num_cells"])
logger = Logger(str(sub_dir), 0, False)
all_archs = meta_info["archs"]
assert srange[1] < meta_info[
"total"], "invalid range : {:}-{:} vs. {:}".format(
srange[0], srange[1], meta_info["total"])
assert (arch_index == -1 or srange[0] <= arch_index <= srange[1]
), "invalid range : {:} vs. {:} vs. {:}".format(
srange[0], arch_index, srange[1])
if arch_index == -1:
to_evaluate_indexes = list(range(srange[0], srange[1] + 1))
else:
to_evaluate_indexes = [arch_index]
logger.log("xargs : seeds = {:}".format(seeds))
logger.log("xargs : arch_index = {:}".format(arch_index))
logger.log("xargs : cover_mode = {:}".format(cover_mode))
logger.log("-" * 100)
logger.log(
"Start evaluating range =: {:06d} vs. {:06d} vs. {:06d} / {:06d} with cover-mode={:}"
.format(srange[0], arch_index, srange[1], meta_info["total"],
cover_mode))
for i, (dataset, xpath, split) in enumerate(zip(datasets, xpaths, splits)):
logger.log(
"--->>> Evaluate {:}/{:} : dataset={:9s}, path={:}, split={:}".
format(i, len(datasets), dataset, xpath, split))
logger.log("--->>> architecture config : {:}".format(arch_config))
start_time, epoch_time = time.time(), AverageMeter()
for i, index in enumerate(to_evaluate_indexes):
arch = all_archs[index]
logger.log(
"\n{:} evaluate {:06d}/{:06d} ({:06d}/{:06d})-th architecture [seeds={:}] {:}"
.format(
"-" * 15,
i,
len(to_evaluate_indexes),
index,
meta_info["total"],
seeds,
"-" * 15,
))
# logger.log('{:} {:} {:}'.format('-'*15, arch.tostr(), '-'*15))
logger.log("{:} {:} {:}".format("-" * 15, arch, "-" * 15))
# test this arch on different datasets with different seeds
has_continue = False
for seed in seeds:
to_save_name = sub_dir / "arch-{:06d}-seed-{:04d}.pth".format(
index, seed)
if to_save_name.exists():
if cover_mode:
logger.log(
"Find existing file : {:}, remove it before evaluation"
.format(to_save_name))
os.remove(str(to_save_name))
else:
logger.log(
"Find existing file : {:}, skip this evaluation".
format(to_save_name))
has_continue = True
continue
results = evaluate_all_datasets(
CellStructure.str2structure(arch),
datasets,
xpaths,
splits,
use_less,
seed,
arch_config,
workers,
logger,
)
torch.save(results, to_save_name)
logger.log(
"{:} --evaluate-- {:06d}/{:06d} ({:06d}/{:06d})-th seed={:} done, save into {:}"
.format(
"-" * 15,
i,
len(to_evaluate_indexes),
index,
meta_info["total"],
seed,
to_save_name,
))
# measure elapsed time
if not has_continue:
epoch_time.update(time.time() - start_time)
start_time = time.time()
need_time = "Time Left: {:}".format(
convert_secs2time(
epoch_time.avg * (len(to_evaluate_indexes) - i - 1), True))
logger.log("This arch costs : {:}".format(
convert_secs2time(epoch_time.val, True)))
logger.log("{:}".format("*" * 100))
logger.log("{:} {:74s} {:}".format(
"*" * 10,
"{:06d}/{:06d} ({:06d}/{:06d})-th done, left {:}".format(
i, len(to_evaluate_indexes), index, meta_info["total"],
need_time),
"*" * 10,
))
logger.log("{:}".format("*" * 100))
logger.close()
def main(args):
logger, env_info, model_kwargs = lfna_setup(args)
# check indexes to be evaluated
to_evaluate_indexes = split_str2indexes(args.srange, env_info["total"],
None)
logger.log("Evaluate {:}, which has {:} timestamps in total.".format(
args.srange, len(to_evaluate_indexes)))
w_container_per_epoch = dict()
per_timestamp_time, start_time = AverageMeter(), time.time()
for i, idx in enumerate(to_evaluate_indexes):
need_time = "Time Left: {:}".format(
convert_secs2time(
per_timestamp_time.avg * (len(to_evaluate_indexes) - i), True))
logger.log("[{:}]".format(time_string()) +
" [{:04d}/{:04d}][{:04d}]".format(i, len(
to_evaluate_indexes), idx) + " " + need_time)
# train the same data
assert idx != 0
historical_x, historical_y = [], []
for past_i in range(idx):
historical_x.append(env_info["{:}-x".format(past_i)])
historical_y.append(env_info["{:}-y".format(past_i)])
historical_x, historical_y = torch.cat(historical_x), torch.cat(
historical_y)
historical_x, historical_y = subsample(historical_x, historical_y)
# build model
model = get_model(dict(model_type="simple_mlp"), **model_kwargs)
# build optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.init_lr,
amsgrad=True)
criterion = torch.nn.MSELoss()
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[
int(args.epochs * 0.25),
int(args.epochs * 0.5),
int(args.epochs * 0.75),
],
gamma=0.3,
)
train_metric = MSEMetric()
best_loss, best_param = None, None
for _iepoch in range(args.epochs):
preds = model(historical_x)
optimizer.zero_grad()
loss = criterion(preds, historical_y)
loss.backward()
optimizer.step()
lr_scheduler.step()
# save best
if best_loss is None or best_loss > loss.item():
best_loss = loss.item()
best_param = copy.deepcopy(model.state_dict())
model.load_state_dict(best_param)
with torch.no_grad():
train_metric(preds, historical_y)
train_results = train_metric.get_info()
metric = ComposeMetric(MSEMetric(), SaveMetric())
eval_dataset = torch.utils.data.TensorDataset(
env_info["{:}-x".format(idx)], env_info["{:}-y".format(idx)])
eval_loader = torch.utils.data.DataLoader(eval_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=0)
results = basic_eval_fn(eval_loader, model, metric, logger)
log_str = ("[{:}]".format(time_string()) +
" [{:04d}/{:04d}]".format(idx, env_info["total"]) +
" train-mse: {:.5f}, eval-mse: {:.5f}".format(
train_results["mse"], results["mse"]))
logger.log(log_str)
save_path = logger.path(None) / "{:04d}-{:04d}.pth".format(
idx, env_info["total"])
w_container_per_epoch[idx] = model.get_w_container().no_grad_clone()
save_checkpoint(
{
"model_state_dict": model.state_dict(),
"model": model,
"index": idx,
"timestamp": env_info["{:}-timestamp".format(idx)],
},
save_path,
logger,
)
logger.log("")
per_timestamp_time.update(time.time() - start_time)
start_time = time.time()
save_checkpoint(
{"w_container_per_epoch": w_container_per_epoch},
logger.path(None) / "final-ckp.pth",
logger,
)
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
)
# 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": "DARTS-V1",
"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))
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)
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,
xargs.gradient_clip,
)
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() ))
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)
logger.log(
"DARTS-V1 : 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, valid_data, xshape, class_num = get_datasets(
xargs.dataset, xargs.data_path, -1
)
if xargs.overwite_epochs is None:
extra_info = {"class_num": class_num, "xshape": xshape}
else:
extra_info = {
"class_num": class_num,
"xshape": xshape,
"epochs": xargs.overwite_epochs,
}
config = load_config(xargs.config_path, extra_info, logger)
search_loader, train_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(xargs.search_space, "nats-bench")
model_config = dict2config(
dict(
name="generic",
super_type="search-shape",
candidate_Cs=search_space["candidates"],
max_num_Cs=search_space["numbers"],
num_classes=class_num,
genotype=args.genotype,
affine=bool(xargs.affine),
track_running_stats=bool(xargs.track_running_stats),
),
None,
)
logger.log("search space : {:}".format(search_space))
logger.log("model config : {:}".format(model_config))
search_model = get_cell_based_tiny_net(model_config)
search_model.set_algo(xargs.algo)
logger.log("{:}".format(search_model))
w_optimizer, w_scheduler, criterion = get_optim_scheduler(
search_model.weights, config
)
a_optimizer = torch.optim.Adam(
search_model.alphas,
lr=xargs.arch_learning_rate,
betas=(0.5, 0.999),
weight_decay=xargs.arch_weight_decay,
eps=xargs.arch_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))
params = count_parameters_in_MB(search_model)
logger.log("The parameters of the search model = {:.2f} MB".format(params))
logger.log("search-space : {:}".format(search_space))
if bool(xargs.use_api):
api = create(None, "size", fast_mode=True, verbose=False)
else:
api = None
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 = search_model.cuda(), criterion.cuda() # use a single GPU
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"]
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: network.random}
# 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)
if (
xargs.warmup_ratio is None
or xargs.warmup_ratio <= float(epoch) / total_epoch
):
enable_controller = True
network.set_warmup_ratio(None)
else:
enable_controller = False
network.set_warmup_ratio(
1.0 - float(epoch) / total_epoch / xargs.warmup_ratio
)
logger.log(
"\n[Search the {:}-th epoch] {:}, LR={:}, controller-warmup={:}, enable_controller={:}".format(
epoch_str,
need_time,
min(w_scheduler.get_lr()),
network.warmup_ratio,
enable_controller,
)
)
if xargs.algo == "mask_gumbel" or xargs.algo == "tas":
network.set_tau(
xargs.tau_max
- (xargs.tau_max - xargs.tau_min) * epoch / (total_epoch - 1)
)
logger.log("[RESET tau as : {:}]".format(network.tau))
(
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,
enable_controller,
xargs.algo,
epoch_str,
xargs.print_freq,
logger,
)
search_time.update(time.time() - start_time)
logger.log(
"[{:}] search [base] : 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(
"[{:}] search [arch] : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%".format(
epoch_str, search_a_loss, search_a_top1, search_a_top5
)
)
genotype = network.genotype
logger.log("[{:}] - [get_best_arch] : {:}".format(epoch_str, genotype))
valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
valid_loader, network, criterion, logger
)
logger.log(
"[{:}] evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}% | {:}".format(
epoch_str, valid_a_loss, valid_a_top1, valid_a_top5, genotype
)
)
valid_accuracies[epoch] = valid_a_top1
genotypes[epoch] = 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,
)
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], "90")))
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
# the final post procedure : count the time
start_time = time.time()
genotype = network.genotype
search_time.update(time.time() - start_time)
valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
valid_loader, network, criterion, logger
)
logger.log(
"Last : the gentotype is : {:}, with the validation accuracy of {:.3f}%.".format(
genotype, valid_a_top1
)
)
logger.log("\n" + "-" * 100)
# check the performance from the architecture dataset
logger.log(
"[{:}] run {:} epochs, cost {:.1f} s, last-geno is {:}.".format(
xargs.algo, total_epoch, search_time.sum, genotype
)
)
if api is not None:
logger.log("{:}".format(api.query_by_arch(genotype, "90")))
logger.close()
def search_func(
xloader,
network,
criterion,
scheduler,
w_optimizer,
a_optimizer,
enable_controller,
algo,
epoch_str,
print_freq,
logger,
):
data_time, batch_time = AverageMeter(), AverageMeter()
base_losses, base_top1, base_top5 = AverageMeter(), AverageMeter(), AverageMeter()
arch_losses, arch_top1, arch_top5 = AverageMeter(), AverageMeter(), AverageMeter()
end = time.time()
network.train()
for step, (base_inputs, base_targets, arch_inputs, arch_targets) in enumerate(
xloader
):
scheduler.update(None, 1.0 * step / len(xloader))
base_inputs = base_inputs.cuda(non_blocking=True)
arch_inputs = arch_inputs.cuda(non_blocking=True)
base_targets = base_targets.cuda(non_blocking=True)
arch_targets = arch_targets.cuda(non_blocking=True)
# measure data loading time
data_time.update(time.time() - end)
# Update the weights
network.zero_grad()
_, logits, _ = network(base_inputs)
base_loss = criterion(logits, base_targets)
base_loss.backward()
w_optimizer.step()
# record
base_prec1, base_prec5 = obtain_accuracy(
logits.data, base_targets.data, topk=(1, 5)
)
base_losses.update(base_loss.item(), base_inputs.size(0))
base_top1.update(base_prec1.item(), base_inputs.size(0))
base_top5.update(base_prec5.item(), base_inputs.size(0))
# update the architecture-weight
network.zero_grad()
a_optimizer.zero_grad()
_, logits, log_probs = network(arch_inputs)
arch_prec1, arch_prec5 = obtain_accuracy(
logits.data, arch_targets.data, topk=(1, 5)
)
if algo == "mask_rl":
with torch.no_grad():
RL_BASELINE_EMA.update(arch_prec1.item())
rl_advantage = arch_prec1 - RL_BASELINE_EMA.value
rl_log_prob = sum(log_probs)
arch_loss = -rl_advantage * rl_log_prob
elif algo == "tas" or algo == "mask_gumbel":
arch_loss = criterion(logits, arch_targets)
else:
raise ValueError("invalid algorightm name: {:}".format(algo))
if enable_controller:
arch_loss.backward()
a_optimizer.step()
# record
arch_losses.update(arch_loss.item(), arch_inputs.size(0))
arch_top1.update(arch_prec1.item(), arch_inputs.size(0))
arch_top5.update(arch_prec5.item(), arch_inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if step % print_freq == 0 or step + 1 == len(xloader):
Sstr = (
"*SEARCH* "
+ time_string()
+ " [{:}][{:03d}/{:03d}]".format(epoch_str, step, len(xloader))
)
Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format(
batch_time=batch_time, data_time=data_time
)
Wstr = "Base [Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})]".format(
loss=base_losses, top1=base_top1, top5=base_top5
)
Astr = "Arch [Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})]".format(
loss=arch_losses, top1=arch_top1, top5=arch_top5
)
logger.log(Sstr + " " + Tstr + " " + Wstr + " " + Astr)
return (
base_losses.avg,
base_top1.avg,
base_top5.avg,
arch_losses.avg,
arch_top1.avg,
arch_top5.avg,
)
def main(args):
prepare_seed(args.rand_seed)
logger = prepare_logger(args)
env = get_synthetic_env(mode=None, version=args.env_version)
model_kwargs = dict(
config=dict(model_type="norm_mlp"),
input_dim=env.meta_info["input_dim"],
output_dim=env.meta_info["output_dim"],
hidden_dims=[args.hidden_dim] * 2,
act_cls="relu",
norm_cls="layer_norm_1d",
)
logger.log("The total enviornment: {:}".format(env))
w_containers = dict()
if env.meta_info["task"] == "regression":
criterion = torch.nn.MSELoss()
metric_cls = MSEMetric
elif env.meta_info["task"] == "classification":
criterion = torch.nn.CrossEntropyLoss()
metric_cls = Top1AccMetric
else:
raise ValueError("This task ({:}) is not supported.".format(
all_env.meta_info["task"]))
per_timestamp_time, start_time = AverageMeter(), time.time()
for idx, (future_time, (future_x, future_y)) in enumerate(env):
need_time = "Time Left: {:}".format(
convert_secs2time(per_timestamp_time.avg * (len(env) - idx), True))
logger.log("[{:}]".format(time_string()) +
" [{:04d}/{:04d}]".format(idx, len(env)) + " " + need_time)
# train the same data
historical_x = future_x.to(args.device)
historical_y = future_y.to(args.device)
# build model
model = get_model(**model_kwargs)
model = model.to(args.device)
if idx == 0:
print(model)
# build optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.init_lr,
amsgrad=True)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[
int(args.epochs * 0.25),
int(args.epochs * 0.5),
int(args.epochs * 0.75),
],
gamma=0.3,
)
train_metric = metric_cls(True)
best_loss, best_param = None, None
for _iepoch in range(args.epochs):
preds = model(historical_x)
optimizer.zero_grad()
loss = criterion(preds, historical_y)
loss.backward()
optimizer.step()
lr_scheduler.step()
# save best
if best_loss is None or best_loss > loss.item():
best_loss = loss.item()
best_param = copy.deepcopy(model.state_dict())
model.load_state_dict(best_param)
model.analyze_weights()
with torch.no_grad():
train_metric(preds, historical_y)
train_results = train_metric.get_info()
xmetric = ComposeMetric(metric_cls(True), SaveMetric())
eval_dataset = torch.utils.data.TensorDataset(future_x.to(args.device),
future_y.to(args.device))
eval_loader = torch.utils.data.DataLoader(eval_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=0)
results = basic_eval_fn(eval_loader, model, xmetric, logger)
log_str = ("[{:}]".format(time_string()) +
" [{:04d}/{:04d}]".format(idx, len(env)) +
" train-score: {:.5f}, eval-score: {:.5f}".format(
train_results["score"], results["score"]))
logger.log(log_str)
save_path = logger.path(None) / "{:04d}-{:04d}.pth".format(
idx, len(env))
w_containers[idx] = model.get_w_container().no_grad_clone()
save_checkpoint(
{
"model_state_dict": model.state_dict(),
"model": model,
"index": idx,
"timestamp": future_time.item(),
},
save_path,
logger,
)
logger.log("")
per_timestamp_time.update(time.time() - start_time)
start_time = time.time()
save_checkpoint(
{"w_containers": w_containers},
logger.path(None) / "final-ckp.pth",
logger,
)
logger.log("-" * 200 + "\n")
logger.close()
def main(args):
logger, model_kwargs = lfna_setup(args)
w_containers = dict()
per_timestamp_time, start_time = AverageMeter(), time.time()
for idx in range(args.prev_time, env_info["total"]):
need_time = "Time Left: {:}".format(
convert_secs2time(per_timestamp_time.avg * (env_info["total"] - idx), True)
)
logger.log(
"[{:}]".format(time_string())
+ " [{:04d}/{:04d}]".format(idx, env_info["total"])
+ " "
+ need_time
)
# train the same data
historical_x = env_info["{:}-x".format(idx - args.prev_time)]
historical_y = env_info["{:}-y".format(idx - args.prev_time)]
# build model
model = get_model(**model_kwargs)
print(model)
# build optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.init_lr, amsgrad=True)
criterion = torch.nn.MSELoss()
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[
int(args.epochs * 0.25),
int(args.epochs * 0.5),
int(args.epochs * 0.75),
],
gamma=0.3,
)
train_metric = MSEMetric()
best_loss, best_param = None, None
for _iepoch in range(args.epochs):
preds = model(historical_x)
optimizer.zero_grad()
loss = criterion(preds, historical_y)
loss.backward()
optimizer.step()
lr_scheduler.step()
# save best
if best_loss is None or best_loss > loss.item():
best_loss = loss.item()
best_param = copy.deepcopy(model.state_dict())
model.load_state_dict(best_param)
model.analyze_weights()
with torch.no_grad():
train_metric(preds, historical_y)
train_results = train_metric.get_info()
metric = ComposeMetric(MSEMetric(), SaveMetric())
eval_dataset = torch.utils.data.TensorDataset(
env_info["{:}-x".format(idx)], env_info["{:}-y".format(idx)]
)
eval_loader = torch.utils.data.DataLoader(
eval_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0
)
results = basic_eval_fn(eval_loader, model, metric, logger)
log_str = (
"[{:}]".format(time_string())
+ " [{:04d}/{:04d}]".format(idx, env_info["total"])
+ " train-mse: {:.5f}, eval-mse: {:.5f}".format(
train_results["mse"], results["mse"]
)
)
logger.log(log_str)
save_path = logger.path(None) / "{:04d}-{:04d}.pth".format(
idx, env_info["total"]
)
w_containers[idx] = model.get_w_container().no_grad_clone()
save_checkpoint(
{
"model_state_dict": model.state_dict(),
"model": model,
"index": idx,
"timestamp": env_info["{:}-timestamp".format(idx)],
},
save_path,
logger,
)
logger.log("")
per_timestamp_time.update(time.time() - start_time)
start_time = time.time()
save_checkpoint(
{"w_containers": w_containers},
logger.path(None) / "final-ckp.pth",
logger,
)
logger.log("-" * 200 + "\n")
logger.close()
def train_single_model(save_dir, workers, datasets, xpaths, splits, use_less,
seeds, model_str, arch_config):
assert torch.cuda.is_available(), "CUDA is not available."
torch.backends.cudnn.enabled = True
torch.backends.cudnn.deterministic = True
# torch.backends.cudnn.benchmark = True
torch.set_num_threads(workers)
save_dir = (Path(save_dir) / "specifics" / "{:}-{:}-{:}-{:}".format(
"LESS" if use_less else "FULL",
model_str,
arch_config["channel"],
arch_config["num_cells"],
))
logger = Logger(str(save_dir), 0, False)
if model_str in CellArchitectures:
arch = CellArchitectures[model_str]
logger.log(
"The model string is found in pre-defined architecture dict : {:}".
format(model_str))
else:
try:
arch = CellStructure.str2structure(model_str)
except:
raise ValueError(
"Invalid model string : {:}. It can not be found or parsed.".
format(model_str))
assert arch.check_valid_op(get_search_spaces(
"cell", "full")), "{:} has the invalid op.".format(arch)
logger.log("Start train-evaluate {:}".format(arch.tostr()))
logger.log("arch_config : {:}".format(arch_config))
start_time, seed_time = time.time(), AverageMeter()
for _is, seed in enumerate(seeds):
logger.log(
"\nThe {:02d}/{:02d}-th seed is {:} ----------------------<.>----------------------"
.format(_is, len(seeds), seed))
to_save_name = save_dir / "seed-{:04d}.pth".format(seed)
if to_save_name.exists():
logger.log("Find the existing file {:}, directly load!".format(
to_save_name))
checkpoint = torch.load(to_save_name)
else:
logger.log(
"Does not find the existing file {:}, train and evaluate!".
format(to_save_name))
checkpoint = evaluate_all_datasets(
arch,
datasets,
xpaths,
splits,
use_less,
seed,
arch_config,
workers,
logger,
)
torch.save(checkpoint, to_save_name)
# log information
logger.log("{:}".format(checkpoint["info"]))
all_dataset_keys = checkpoint["all_dataset_keys"]
for dataset_key in all_dataset_keys:
logger.log("\n{:} dataset : {:} {:}".format(
"-" * 15, dataset_key, "-" * 15))
dataset_info = checkpoint[dataset_key]
# logger.log('Network ==>\n{:}'.format( dataset_info['net_string'] ))
logger.log("Flops = {:} MB, Params = {:} MB".format(
dataset_info["flop"], dataset_info["param"]))
logger.log("config : {:}".format(dataset_info["config"]))
logger.log("Training State (finish) = {:}".format(
dataset_info["finish-train"]))
last_epoch = dataset_info["total_epoch"] - 1
train_acc1es, train_acc5es = (
dataset_info["train_acc1es"],
dataset_info["train_acc5es"],
)
valid_acc1es, valid_acc5es = (
dataset_info["valid_acc1es"],
dataset_info["valid_acc5es"],
)
logger.log(
"Last Info : Train = Acc@1 {:.2f}% Acc@5 {:.2f}% Error@1 {:.2f}%, Test = Acc@1 {:.2f}% Acc@5 {:.2f}% Error@1 {:.2f}%"
.format(
train_acc1es[last_epoch],
train_acc5es[last_epoch],
100 - train_acc1es[last_epoch],
valid_acc1es[last_epoch],
valid_acc5es[last_epoch],
100 - valid_acc1es[last_epoch],
))
# measure elapsed time
seed_time.update(time.time() - start_time)
start_time = time.time()
need_time = "Time Left: {:}".format(
convert_secs2time(seed_time.avg * (len(seeds) - _is - 1), True))
logger.log(
"\n<<<***>>> The {:02d}/{:02d}-th seed is {:} <finish> other procedures need {:}"
.format(_is, len(seeds), seed, need_time))
logger.close()
def meta_train_procedure(base_model, meta_model, criterion, xenv, args,
logger):
base_model.train()
meta_model.train()
optimizer = torch.optim.Adam(
meta_model.get_parameters(True, True, True),
lr=args.lr,
weight_decay=args.weight_decay,
amsgrad=True,
)
logger.log("Pre-train the meta-model")
logger.log("Using the optimizer: {:}".format(optimizer))
meta_model.set_best_dir(logger.path(None) / "ckps-pretrain-v2")
final_best_name = "final-pretrain-{:}.pth".format(args.rand_seed)
if meta_model.has_best(final_best_name):
meta_model.load_best(final_best_name)
logger.log(
"Directly load the best model from {:}".format(final_best_name))
return
total_indexes = list(range(meta_model.meta_length))
meta_model.set_best_name("pretrain-{:}.pth".format(args.rand_seed))
last_success_epoch, early_stop_thresh = 0, args.pretrain_early_stop_thresh
per_epoch_time, start_time = AverageMeter(), time.time()
device = args.device
for iepoch in range(args.epochs):
left_time = "Time Left: {:}".format(
convert_secs2time(per_epoch_time.avg * (args.epochs - iepoch),
True))
optimizer.zero_grad()
generated_time_embeds = meta_model.gen_time_embed(
meta_model.meta_timestamps)
batch_indexes = random.choices(total_indexes, k=args.meta_batch)
raw_time_steps = meta_model.meta_timestamps[batch_indexes]
regularization_loss = F.l1_loss(generated_time_embeds,
meta_model.super_meta_embed,
reduction="mean")
# future loss
total_future_losses, total_present_losses = [], []
future_containers = meta_model.gen_model(
generated_time_embeds[batch_indexes])
present_containers = meta_model.gen_model(
meta_model.super_meta_embed[batch_indexes])
for ibatch, time_step in enumerate(raw_time_steps.cpu().tolist()):
_, (inputs, targets) = xenv(time_step)
inputs, targets = inputs.to(device), targets.to(device)
predictions = base_model.forward_with_container(
inputs, future_containers[ibatch])
total_future_losses.append(criterion(predictions, targets))
predictions = base_model.forward_with_container(
inputs, present_containers[ibatch])
total_present_losses.append(criterion(predictions, targets))
with torch.no_grad():
meta_std = torch.stack(total_future_losses).std().item()
loss_future = torch.stack(total_future_losses).mean()
loss_present = torch.stack(total_present_losses).mean()
total_loss = loss_future + loss_present + regularization_loss
total_loss.backward()
optimizer.step()
# success
success, best_score = meta_model.save_best(-total_loss.item())
logger.log(
"{:} [META {:04d}/{:}] loss : {:.4f} +- {:.4f} = {:.4f} + {:.4f} + {:.4f}"
.format(
time_string(),
iepoch,
args.epochs,
total_loss.item(),
meta_std,
loss_future.item(),
loss_present.item(),
regularization_loss.item(),
) + ", batch={:}".format(len(total_future_losses)) +
", success={:}, best={:.4f}".format(success, -best_score) +
", LS={:}/{:}".format(iepoch -
last_success_epoch, early_stop_thresh) +
", {:}".format(left_time))
if success:
last_success_epoch = iepoch
if iepoch - last_success_epoch >= early_stop_thresh:
logger.log("Early stop the pre-training at {:}".format(iepoch))
break
per_epoch_time.update(time.time() - start_time)
start_time = time.time()
meta_model.load_best()
# save to the final model
meta_model.set_best_name(final_best_name)
success, _ = meta_model.save_best(best_score + 1e-6)
assert success
logger.log("Save the best model into {:}".format(final_best_name))
def main(
save_dir: Path,
workers: int,
datasets: List[Text],
xpaths: List[Text],
splits: List[int],
seeds: List[int],
nets: List[str],
opt_config: Dict[Text, Any],
to_evaluate_indexes: tuple,
cover_mode: bool,
arch_config: Dict[Text, Any],
):
log_dir = save_dir / "logs"
log_dir.mkdir(parents=True, exist_ok=True)
logger = Logger(str(log_dir), os.getpid(), False)
logger.log("xargs : seeds = {:}".format(seeds))
logger.log("xargs : cover_mode = {:}".format(cover_mode))
logger.log("-" * 100)
logger.log("Start evaluating range =: {:06d} - {:06d}".format(
min(to_evaluate_indexes), max(to_evaluate_indexes)) +
"({:} in total) / {:06d} with cover-mode={:}".format(
len(to_evaluate_indexes), len(nets), cover_mode))
for i, (dataset, xpath, split) in enumerate(zip(datasets, xpaths, splits)):
logger.log(
"--->>> Evaluate {:}/{:} : dataset={:9s}, path={:}, split={:}".
format(i, len(datasets), dataset, xpath, split))
logger.log("--->>> optimization config : {:}".format(opt_config))
start_time, epoch_time = time.time(), AverageMeter()
for i, index in enumerate(to_evaluate_indexes):
arch = nets[index]
logger.log(
"\n{:} evaluate {:06d}/{:06d} ({:06d}/{:06d})-th arch [seeds={:}] {:}"
.format(
time_string(),
i,
len(to_evaluate_indexes),
index,
len(nets),
seeds,
"-" * 15,
))
logger.log("{:} {:} {:}".format("-" * 15, arch, "-" * 15))
# test this arch on different datasets with different seeds
has_continue = False
for seed in seeds:
to_save_name = save_dir / "arch-{:06d}-seed-{:04d}.pth".format(
index, seed)
if to_save_name.exists():
if cover_mode:
logger.log(
"Find existing file : {:}, remove it before evaluation"
.format(to_save_name))
os.remove(str(to_save_name))
else:
logger.log(
"Find existing file : {:}, skip this evaluation".
format(to_save_name))
has_continue = True
continue
results = evaluate_all_datasets(
CellStructure.str2structure(arch),
datasets,
xpaths,
splits,
opt_config,
seed,
arch_config,
workers,
logger,
)
torch.save(results, to_save_name)
logger.log(
"\n{:} evaluate {:06d}/{:06d} ({:06d}/{:06d})-th arch [seeds={:}] ===>>> {:}"
.format(
time_string(),
i,
len(to_evaluate_indexes),
index,
len(nets),
seeds,
to_save_name,
))
# measure elapsed time
if not has_continue:
epoch_time.update(time.time() - start_time)
start_time = time.time()
need_time = "Time Left: {:}".format(
convert_secs2time(
epoch_time.avg * (len(to_evaluate_indexes) - i - 1), True))
logger.log("This arch costs : {:}".format(
convert_secs2time(epoch_time.val, True)))
logger.log("{:}".format("*" * 100))
logger.log("{:} {:74s} {:}".format(
"*" * 10,
"{:06d}/{:06d} ({:06d}/{:06d})-th done, left {:}".format(
i, len(to_evaluate_indexes), index, len(nets), need_time),
"*" * 10,
))
logger.log("{:}".format("*" * 100))
logger.close()
def procedure(xloader, network, criterion, scheduler, optimizer, mode, config,
extra_info, print_freq, logger, local_epoch):
data_time, batch_time, losses, top1, top5 = (
AverageMeter(),
AverageMeter(),
AverageMeter(),
AverageMeter(),
AverageMeter(),
)
if mode == "train":
network.train()
elif mode == "valid":
network.eval()
else:
raise ValueError("The mode is not right : {:}".format(mode))
# logger.log('[{:5s}] config :: auxiliary={:}, message={:}'.format(mode, config.auxiliary if hasattr(config, 'auxiliary') else -1, network.module.get_message()))
# logger.log(
# "[{:5s}] config :: auxiliary={:}".format(
# mode, config.auxiliary if hasattr(config, "auxiliary") else -1
# )
# )
end = time.time()
loss_aux = None
for epoch in range(local_epoch):
for i, (inputs, targets) in enumerate(xloader):
if mode == "train":
scheduler.update(None, 1.0 * i / len(xloader))
# measure data loading time
data_time.update(time.time() - end)
# calculate prediction and loss
targets = targets.cuda(non_blocking=True)
inputs = inputs.to('cuda')
if mode == "train":
optimizer.zero_grad()
logits = network(inputs)
if isinstance(logits, tuple):
features, logits = logits
if isinstance(logits, list):
assert len(
logits
) == 2, "logits must has {:} items instead of {:}".format(
2, len(logits))
logits, logits_aux = logits
else:
logits, logits_aux = logits, None
loss = criterion(logits, targets)
# if loss_aux is not None and config is not None and hasattr(config, "auxiliary") and config.auxiliary > 0:
# loss_aux = criterion(logits_aux, targets)
# loss += config.auxiliary * loss_aux
if mode == "train":
loss.backward()
optimizer.step()
# record
prec1, prec5 = obtain_accuracy(logits.data,
targets.data,
topk=(1, 5))
losses.update(loss.item(), inputs.size(0))
top1.update(prec1.item(), inputs.size(0))
top5.update(prec5.item(), inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % print_freq == 0 or (i + 1) == len(xloader):
Sstr = (" {:5s} ".format(mode.upper()) + time_string() +
" [{:}][{:03d}/{:03d}]".format(extra_info, i,
len(xloader)))
if scheduler is not None:
Sstr += " {:}".format(scheduler.get_min_info())
Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format(
batch_time=batch_time, data_time=data_time)
Lstr = "Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})".format(
loss=losses, top1=top1, top5=top5)
Istr = "Size={:}".format(list(inputs.size()))
logger.log(Sstr + " " + Tstr + " " + Lstr + " " + Istr)
logger.log(
" **{mode:5s}** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Loss:{loss:.3f}"
.format(
mode=mode.upper(),
top1=top1,
top5=top5,
error1=100 - top1.avg,
error5=100 - top5.avg,
loss=losses.avg,
))
return losses.avg, top1.avg, top5.avg
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(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(args):
prepare_seed(args.rand_seed)
logger = prepare_logger(args)
env = get_synthetic_env(mode="test", version=args.env_version)
model_kwargs = dict(
config=dict(model_type="norm_mlp"),
input_dim=env.meta_info["input_dim"],
output_dim=env.meta_info["output_dim"],
hidden_dims=[args.hidden_dim] * 2,
act_cls="relu",
norm_cls="layer_norm_1d",
)
logger.log("The total enviornment: {:}".format(env))
w_containers = dict()
if env.meta_info["task"] == "regression":
criterion = torch.nn.MSELoss()
metric_cls = MSEMetric
elif env.meta_info["task"] == "classification":
criterion = torch.nn.CrossEntropyLoss()
metric_cls = Top1AccMetric
else:
raise ValueError("This task ({:}) is not supported.".format(
all_env.meta_info["task"]))
def finetune(index):
seq_times = env.get_seq_times(index, args.seq_length)
_, (allxs, allys) = env.seq_call(seq_times)
allxs, allys = allxs.view(-1, allxs.shape[-1]), allys.view(-1, 1)
if env.meta_info["task"] == "classification":
allys = allys.view(-1)
historical_x, historical_y = allxs.to(args.device), allys.to(
args.device)
model = get_model(**model_kwargs)
model = model.to(args.device)
optimizer = torch.optim.Adam(model.parameters(),
lr=args.init_lr,
amsgrad=True)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[
int(args.epochs * 0.25),
int(args.epochs * 0.5),
int(args.epochs * 0.75),
],
gamma=0.3,
)
train_metric = metric_cls(True)
best_loss, best_param = None, None
for _iepoch in range(args.epochs):
preds = model(historical_x)
optimizer.zero_grad()
loss = criterion(preds, historical_y)
loss.backward()
optimizer.step()
lr_scheduler.step()
# save best
if best_loss is None or best_loss > loss.item():
best_loss = loss.item()
best_param = copy.deepcopy(model.state_dict())
model.load_state_dict(best_param)
# model.analyze_weights()
with torch.no_grad():
train_metric(preds, historical_y)
train_results = train_metric.get_info()
return train_results, model
metric = metric_cls(True)
per_timestamp_time, start_time = AverageMeter(), time.time()
for idx, (future_time, (future_x, future_y)) in enumerate(env):
need_time = "Time Left: {:}".format(
convert_secs2time(per_timestamp_time.avg * (len(env) - idx), True))
logger.log("[{:}]".format(time_string()) +
" [{:04d}/{:04d}]".format(idx, len(env)) + " " + need_time)
# train the same data
train_results, model = finetune(idx)
# build optimizer
xmetric = ComposeMetric(metric_cls(True), SaveMetric())
future_x, future_y = future_x.to(args.device), future_y.to(args.device)
future_y_hat = model(future_x)
future_loss = criterion(future_y_hat, future_y)
metric(future_y_hat, future_y)
log_str = ("[{:}]".format(time_string()) +
" [{:04d}/{:04d}]".format(idx, len(env)) +
" train-score: {:.5f}, eval-score: {:.5f}".format(
train_results["score"],
metric.get_info()["score"]))
logger.log(log_str)
logger.log("")
per_timestamp_time.update(time.time() - start_time)
start_time = time.time()
save_checkpoint(
{"w_containers": w_containers},
logger.path(None) / "final-ckp.pth",
logger,
)
logger.log("-" * 200 + "\n")
logger.close()
return metric.get_info()["score"]
def valid_func(xloader, network, criterion):
data_time, batch_time = AverageMeter(), AverageMeter()
arch_losses, arch_top1, arch_top5 = AverageMeter(), AverageMeter(), AverageMeter()
network.eval()
end = time.time()
with torch.no_grad():
for step, (arch_inputs, arch_targets) in enumerate(xloader):
arch_targets = arch_targets.cuda(non_blocking=True)
# measure data loading time
data_time.update(time.time() - end)
# prediction
_, logits = network(arch_inputs)
arch_loss = criterion(logits, arch_targets)
# record
arch_prec1, arch_prec5 = obtain_accuracy(
logits.data, arch_targets.data, topk=(1, 5)
)
arch_losses.update(arch_loss.item(), arch_inputs.size(0))
arch_top1.update(arch_prec1.item(), arch_inputs.size(0))
arch_top5.update(arch_prec5.item(), arch_inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
return arch_losses.avg, arch_top1.avg, arch_top5.avg
def procedure(
xloader,
teacher,
network,
criterion,
scheduler,
optimizer,
mode,
config,
extra_info,
print_freq,
logger,
):
data_time, batch_time, losses, top1, top5 = (
AverageMeter(),
AverageMeter(),
AverageMeter(),
AverageMeter(),
AverageMeter(),
)
Ttop1, Ttop5 = AverageMeter(), AverageMeter()
if mode == "train":
network.train()
elif mode == "valid":
network.eval()
else:
raise ValueError("The mode is not right : {:}".format(mode))
teacher.eval()
logger.log(
"[{:5s}] config :: auxiliary={:}, KD :: [alpha={:.2f}, temperature={:.2f}]"
.format(
mode,
config.auxiliary if hasattr(config, "auxiliary") else -1,
config.KD_alpha,
config.KD_temperature,
))
end = time.time()
for i, (inputs, targets) in enumerate(xloader):
if mode == "train":
scheduler.update(None, 1.0 * i / len(xloader))
# measure data loading time
data_time.update(time.time() - end)
# calculate prediction and loss
targets = targets.cuda(non_blocking=True)
if mode == "train":
optimizer.zero_grad()
student_f, logits = network(inputs)
if isinstance(logits, list):
assert len(
logits
) == 2, "logits must has {:} items instead of {:}".format(
2, len(logits))
logits, logits_aux = logits
else:
logits, logits_aux = logits, None
with torch.no_grad():
teacher_f, teacher_logits = teacher(inputs)
loss = loss_KD_fn(
criterion,
logits,
teacher_logits,
student_f,
teacher_f,
targets,
config.KD_alpha,
config.KD_temperature,
)
if config is not None and hasattr(
config, "auxiliary") and config.auxiliary > 0:
loss_aux = criterion(logits_aux, targets)
loss += config.auxiliary * loss_aux
if mode == "train":
loss.backward()
optimizer.step()
# record
sprec1, sprec5 = obtain_accuracy(logits.data,
targets.data,
topk=(1, 5))
losses.update(loss.item(), inputs.size(0))
top1.update(sprec1.item(), inputs.size(0))
top5.update(sprec5.item(), inputs.size(0))
# teacher
tprec1, tprec5 = obtain_accuracy(teacher_logits.data,
targets.data,
topk=(1, 5))
Ttop1.update(tprec1.item(), inputs.size(0))
Ttop5.update(tprec5.item(), inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % print_freq == 0 or (i + 1) == len(xloader):
Sstr = (
" {:5s} ".format(mode.upper()) + time_string() +
" [{:}][{:03d}/{:03d}]".format(extra_info, i, len(xloader)))
if scheduler is not None:
Sstr += " {:}".format(scheduler.get_min_info())
Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format(
batch_time=batch_time, data_time=data_time)
Lstr = "Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})".format(
loss=losses, top1=top1, top5=top5)
Lstr += " Teacher : acc@1={:.2f}, acc@5={:.2f}".format(
Ttop1.avg, Ttop5.avg)
Istr = "Size={:}".format(list(inputs.size()))
logger.log(Sstr + " " + Tstr + " " + Lstr + " " + Istr)
logger.log(" **{:5s}** accuracy drop :: @1={:.2f}, @5={:.2f}".format(
mode.upper(), Ttop1.avg - top1.avg, Ttop5.avg - top5.avg))
logger.log(
" **{mode:5s}** Prec@1 {top1.avg:.2f} Prec@5 {top5.avg:.2f} Error@1 {error1:.2f} Error@5 {error5:.2f} Loss:{loss:.3f}"
.format(
mode=mode.upper(),
top1=top1,
top5=top5,
error1=100 - top1.avg,
error5=100 - top5.avg,
loss=losses.avg,
))
return losses.avg, top1.avg, top5.avg
def search_func(
xloader,
network,
criterion,
scheduler,
w_optimizer,
a_optimizer,
epoch_str,
print_freq,
logger,
gradient_clip,
):
data_time, batch_time = AverageMeter(), AverageMeter()
base_losses, base_top1, base_top5 = AverageMeter(), AverageMeter(), AverageMeter()
arch_losses, arch_top1, arch_top5 = AverageMeter(), AverageMeter(), AverageMeter()
network.train()
end = time.time()
for step, (base_inputs, base_targets, arch_inputs, arch_targets) in enumerate(
xloader
):
scheduler.update(None, 1.0 * step / len(xloader))
base_targets = base_targets.cuda(non_blocking=True)
arch_targets = arch_targets.cuda(non_blocking=True)
# measure data loading time
data_time.update(time.time() - end)
# update the weights
w_optimizer.zero_grad()
_, logits = network(base_inputs)
base_loss = criterion(logits, base_targets)
base_loss.backward()
if gradient_clip > 0:
torch.nn.utils.clip_grad_norm_(network.parameters(), gradient_clip)
w_optimizer.step()
# record
base_prec1, base_prec5 = obtain_accuracy(
logits.data, base_targets.data, topk=(1, 5)
)
base_losses.update(base_loss.item(), base_inputs.size(0))
base_top1.update(base_prec1.item(), base_inputs.size(0))
base_top5.update(base_prec5.item(), base_inputs.size(0))
# update the architecture-weight
a_optimizer.zero_grad()
_, logits = network(arch_inputs)
arch_loss = criterion(logits, arch_targets)
arch_loss.backward()
a_optimizer.step()
# record
arch_prec1, arch_prec5 = obtain_accuracy(
logits.data, arch_targets.data, topk=(1, 5)
)
arch_losses.update(arch_loss.item(), arch_inputs.size(0))
arch_top1.update(arch_prec1.item(), arch_inputs.size(0))
arch_top5.update(arch_prec5.item(), arch_inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if step % print_freq == 0 or step + 1 == len(xloader):
Sstr = (
"*SEARCH* "
+ time_string()
+ " [{:}][{:03d}/{:03d}]".format(epoch_str, step, len(xloader))
)
Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format(
batch_time=batch_time, data_time=data_time
)
Wstr = "Base [Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})]".format(
loss=base_losses, top1=base_top1, top5=base_top5
)
Astr = "Arch [Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})]".format(
loss=arch_losses, top1=arch_top1, top5=arch_top5
)
logger.log(Sstr + " " + Tstr + " " + Wstr + " " + Astr)
return base_losses.avg, base_top1.avg, base_top5.avg
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()