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):
assert torch.cuda.is_available(), "CUDA is not available."
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
# torch.set_num_threads(args.workers)
prepare_seed(args.rand_seed)
logger = prepare_logger(args)
train_data, valid_data, xshape, class_num = get_datasets(
args.dataset, args.data_path, args.cutout_length)
train_loader = torch.utils.data.DataLoader(
train_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
)
valid_loader = torch.utils.data.DataLoader(
valid_data,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
)
# get configures
model_config = load_config(args.model_config, {"class_num": class_num},
logger)
optim_config = load_config(
args.optim_config,
{
"class_num": class_num,
"KD_alpha": args.KD_alpha,
"KD_temperature": args.KD_temperature,
},
logger,
)
# load checkpoint
teacher_base = load_net_from_checkpoint(args.KD_checkpoint)
teacher = torch.nn.DataParallel(teacher_base).cuda()
base_model = obtain_model(model_config)
flop, param = get_model_infos(base_model, xshape)
logger.log("Student ====>>>>:\n{:}".format(base_model))
logger.log("Teacher ====>>>>:\n{:}".format(teacher_base))
logger.log("model information : {:}".format(base_model.get_message()))
logger.log("-" * 50)
logger.log("Params={:.2f} MB, FLOPs={:.2f} M ... = {:.2f} G".format(
param, flop, flop / 1e3))
logger.log("-" * 50)
logger.log("train_data : {:}".format(train_data))
logger.log("valid_data : {:}".format(valid_data))
optimizer, scheduler, criterion = get_optim_scheduler(
base_model.parameters(), optim_config)
logger.log("optimizer : {:}".format(optimizer))
logger.log("scheduler : {:}".format(scheduler))
logger.log("criterion : {:}".format(criterion))
last_info, model_base_path, model_best_path = (
logger.path("info"),
logger.path("model"),
logger.path("best"),
)
network, criterion = torch.nn.DataParallel(
base_model).cuda(), criterion.cuda()
if last_info.exists(): # automatically resume from previous checkpoint
logger.log("=> loading checkpoint of the last-info '{:}' start".format(
last_info))
last_info = torch.load(last_info)
start_epoch = last_info["epoch"] + 1
checkpoint = torch.load(last_info["last_checkpoint"])
base_model.load_state_dict(checkpoint["base-model"])
scheduler.load_state_dict(checkpoint["scheduler"])
optimizer.load_state_dict(checkpoint["optimizer"])
valid_accuracies = checkpoint["valid_accuracies"]
max_bytes = checkpoint["max_bytes"]
logger.log(
"=> loading checkpoint of the last-info '{:}' start with {:}-th epoch."
.format(last_info, start_epoch))
elif args.resume is not None:
assert Path(
args.resume).exists(), "Can not find the resume file : {:}".format(
args.resume)
checkpoint = torch.load(args.resume)
start_epoch = checkpoint["epoch"] + 1
base_model.load_state_dict(checkpoint["base-model"])
scheduler.load_state_dict(checkpoint["scheduler"])
optimizer.load_state_dict(checkpoint["optimizer"])
valid_accuracies = checkpoint["valid_accuracies"]
max_bytes = checkpoint["max_bytes"]
logger.log(
"=> loading checkpoint from '{:}' start with {:}-th epoch.".format(
args.resume, start_epoch))
elif args.init_model is not None:
assert Path(args.init_model).exists(
), "Can not find the initialization file : {:}".format(args.init_model)
checkpoint = torch.load(args.init_model)
base_model.load_state_dict(checkpoint["base-model"])
start_epoch, valid_accuracies, max_bytes = 0, {"best": -1}, {}
logger.log("=> initialize the model from {:}".format(args.init_model))
else:
logger.log("=> do not find the last-info file : {:}".format(last_info))
start_epoch, valid_accuracies, max_bytes = 0, {"best": -1}, {}
train_func, valid_func = get_procedures(args.procedure)
total_epoch = optim_config.epochs + optim_config.warmup
# Main Training and Evaluation Loop
start_time = time.time()
epoch_time = AverageMeter()
for epoch in range(start_epoch, total_epoch):
scheduler.update(epoch, 0.0)
need_time = "Time Left: {:}".format(
convert_secs2time(epoch_time.avg * (total_epoch - epoch), True))
epoch_str = "epoch={:03d}/{:03d}".format(epoch, total_epoch)
LRs = scheduler.get_lr()
find_best = False
logger.log(
"\n***{:s}*** start {:s} {:s}, LR=[{:.6f} ~ {:.6f}], scheduler={:}"
.format(time_string(), epoch_str, need_time, min(LRs), max(LRs),
scheduler))
# train for one epoch
train_loss, train_acc1, train_acc5 = train_func(
train_loader,
teacher,
network,
criterion,
scheduler,
optimizer,
optim_config,
epoch_str,
args.print_freq,
logger,
)
# log the results
logger.log(
"***{:s}*** TRAIN [{:}] loss = {:.6f}, accuracy-1 = {:.2f}, accuracy-5 = {:.2f}"
.format(time_string(), epoch_str, train_loss, train_acc1,
train_acc5))
# evaluate the performance
if (epoch % args.eval_frequency == 0) or (epoch + 1 == total_epoch):
logger.log("-" * 150)
valid_loss, valid_acc1, valid_acc5 = valid_func(
valid_loader,
teacher,
network,
criterion,
optim_config,
epoch_str,
args.print_freq_eval,
logger,
)
valid_accuracies[epoch] = valid_acc1
logger.log(
"***{:s}*** VALID [{:}] loss = {:.6f}, accuracy@1 = {:.2f}, accuracy@5 = {:.2f} | Best-Valid-Acc@1={:.2f}, Error@1={:.2f}"
.format(
time_string(),
epoch_str,
valid_loss,
valid_acc1,
valid_acc5,
valid_accuracies["best"],
100 - valid_accuracies["best"],
))
if valid_acc1 > valid_accuracies["best"]:
valid_accuracies["best"] = valid_acc1
find_best = True
logger.log(
"Currently, the best validation accuracy found at {:03d}-epoch :: acc@1={:.2f}, acc@5={:.2f}, error@1={:.2f}, error@5={:.2f}, save into {:}."
.format(
epoch,
valid_acc1,
valid_acc5,
100 - valid_acc1,
100 - valid_acc5,
model_best_path,
))
num_bytes = (torch.cuda.max_memory_cached(
next(network.parameters()).device) * 1.0)
logger.log(
"[GPU-Memory-Usage on {:} is {:} bytes, {:.2f} KB, {:.2f} MB, {:.2f} GB.]"
.format(
next(network.parameters()).device,
int(num_bytes),
num_bytes / 1e3,
num_bytes / 1e6,
num_bytes / 1e9,
))
max_bytes[epoch] = num_bytes
if epoch % 10 == 0:
torch.cuda.empty_cache()
# save checkpoint
save_path = save_checkpoint(
{
"epoch": epoch,
"args": deepcopy(args),
"max_bytes": deepcopy(max_bytes),
"FLOP": flop,
"PARAM": param,
"valid_accuracies": deepcopy(valid_accuracies),
"model-config": model_config._asdict(),
"optim-config": optim_config._asdict(),
"base-model": base_model.state_dict(),
"scheduler": scheduler.state_dict(),
"optimizer": optimizer.state_dict(),
},
model_base_path,
logger,
)
if find_best:
copy_checkpoint(model_base_path, model_best_path, logger)
last_info = save_checkpoint(
{
"epoch": epoch,
"args": deepcopy(args),
"last_checkpoint": save_path,
},
logger.path("info"),
logger,
)
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log("\n" + "-" * 200)
logger.log("||| Params={:.2f} MB, FLOPs={:.2f} M ... = {:.2f} G".format(
param, flop, flop / 1e3))
logger.log(
"Finish training/validation in {:} with Max-GPU-Memory of {:.2f} MB, and save final checkpoint into {:}"
.format(
convert_secs2time(epoch_time.sum, True),
max(v for k, v in max_bytes.items()) / 1e6,
logger.path("info"),
))
logger.log("-" * 200 + "\n")
logger.close()
def main(args):
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(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 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 simplify(save_dir, save_name, nets, total):
hps, seeds = ["01", "12", "90"], set()
for hp in hps:
sub_save_dir = save_dir / "raw-data-{:}".format(hp)
ckps = sorted(list(sub_save_dir.glob("arch-*-seed-*.pth")))
seed2names = defaultdict(list)
for ckp in ckps:
parts = re.split("-|\.", ckp.name)
seed2names[parts[3]].append(ckp.name)
print("DIR : {:}".format(sub_save_dir))
nums = []
for seed, xlist in seed2names.items():
seeds.add(seed)
nums.append(len(xlist))
print(" [seed={:}] there are {:} checkpoints.".format(
seed, len(xlist)))
assert (len(nets) == total ==
max(nums)), "there are some missed files : {:} vs {:}".format(
max(nums), total)
print("{:} start simplify the checkpoint.".format(time_string()))
datasets = ("cifar10-valid", "cifar10", "cifar100", "ImageNet16-120")
# Create the directory to save the processed data
# full_save_dir contains all benchmark files with trained weights.
# simplify_save_dir contains all benchmark files without trained weights.
full_save_dir = save_dir / (save_name + "-FULL")
simple_save_dir = save_dir / (save_name + "-SIMPLIFY")
full_save_dir.mkdir(parents=True, exist_ok=True)
simple_save_dir.mkdir(parents=True, exist_ok=True)
# all data in memory
arch2infos, evaluated_indexes = dict(), set()
end_time, arch_time = time.time(), AverageMeter()
for index in tqdm(range(total)):
arch_str = nets[index]
hp2info = OrderedDict()
full_save_path = full_save_dir / "{:06d}.pickle".format(index)
simple_save_path = simple_save_dir / "{:06d}.pickle".format(index)
for hp in hps:
sub_save_dir = save_dir / "raw-data-{:}".format(hp)
ckps = [
sub_save_dir / "arch-{:06d}-seed-{:}.pth".format(index, seed)
for seed in seeds
]
ckps = [x for x in ckps if x.exists()]
if len(ckps) == 0:
raise ValueError("Invalid data : index={:}, hp={:}".format(
index, hp))
arch_info = account_one_arch(index, arch_str, ckps, datasets)
hp2info[hp] = arch_info
hp2info = correct_time_related_info(hp2info)
evaluated_indexes.add(index)
hp2info["01"].clear_params() # to save some spaces...
to_save_data = OrderedDict({
"01": hp2info["01"].state_dict(),
"12": hp2info["12"].state_dict(),
"90": hp2info["90"].state_dict(),
})
pickle_save(to_save_data, str(full_save_path))
for hp in hps:
hp2info[hp].clear_params()
to_save_data = OrderedDict({
"01": hp2info["01"].state_dict(),
"12": hp2info["12"].state_dict(),
"90": hp2info["90"].state_dict(),
})
pickle_save(to_save_data, str(simple_save_path))
arch2infos[index] = to_save_data
# measure elapsed time
arch_time.update(time.time() - end_time)
end_time = time.time()
need_time = "{:}".format(
convert_secs2time(arch_time.avg * (total - index - 1), True))
# print('{:} {:06d}/{:06d} : still need {:}'.format(time_string(), index, total, need_time))
print("{:} {:} done.".format(time_string(), save_name))
final_infos = {
"meta_archs": nets,
"total_archs": total,
"arch2infos": arch2infos,
"evaluated_indexes": evaluated_indexes,
}
save_file_name = save_dir / "{:}.pickle".format(save_name)
pickle_save(final_infos, str(save_file_name))
# move the benchmark file to a new path
hd5sum = get_md5_file(str(save_file_name) + ".pbz2")
hd5_file_name = save_dir / "{:}-{:}.pickle.pbz2".format(
NATS_SSS_BASE_NAME, hd5sum)
shutil.move(str(save_file_name) + ".pbz2", hd5_file_name)
print("Save {:} / {:} architecture results into {:} -> {:}.".format(
len(evaluated_indexes), total, save_file_name, hd5_file_name))
# move the directory to a new path
hd5_full_save_dir = save_dir / "{:}-{:}-full".format(
NATS_SSS_BASE_NAME, hd5sum)
hd5_simple_save_dir = save_dir / "{:}-{:}-simple".format(
NATS_SSS_BASE_NAME, hd5sum)
shutil.move(full_save_dir, hd5_full_save_dir)
shutil.move(simple_save_dir, hd5_simple_save_dir)
# save the meta information for simple and full
final_infos["arch2infos"] = None
final_infos["evaluated_indexes"] = set()
pickle_save(final_infos, str(hd5_full_save_dir / "meta.pickle"))
pickle_save(final_infos, str(hd5_simple_save_dir / "meta.pickle"))
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 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(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 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 evaluate_for_seed(arch_config, opt_config, train_loader, valid_loaders,
seed: int, logger):
"""A modular function to train and evaluate a single network, using the given random seed and optimization config with the provided loaders."""
prepare_seed(seed) # random seed
net = get_cell_based_tiny_net(arch_config)
# net = TinyNetwork(arch_config['channel'], arch_config['num_cells'], arch, config.class_num)
flop, param = get_model_infos(net, opt_config.xshape)
logger.log("Network : {:}".format(net.get_message()), False)
logger.log(
"{:} Seed-------------------------- {:} --------------------------".
format(time_string(), seed))
logger.log("FLOP = {:} MB, Param = {:} MB".format(flop, param))
# train and valid
optimizer, scheduler, criterion = get_optim_scheduler(
net.parameters(), opt_config)
default_device = torch.cuda.current_device()
network = torch.nn.DataParallel(net,
device_ids=[default_device
]).cuda(device=default_device)
criterion = criterion.cuda(device=default_device)
# start training
start_time, epoch_time, total_epoch = (
time.time(),
AverageMeter(),
opt_config.epochs + opt_config.warmup,
)
(
train_losses,
train_acc1es,
train_acc5es,
valid_losses,
valid_acc1es,
valid_acc5es,
) = ({}, {}, {}, {}, {}, {})
train_times, valid_times, lrs = {}, {}, {}
for epoch in range(total_epoch):
scheduler.update(epoch, 0.0)
lr = min(scheduler.get_lr())
train_loss, train_acc1, train_acc5, train_tm = procedure(
train_loader, network, criterion, scheduler, optimizer, "train")
train_losses[epoch] = train_loss
train_acc1es[epoch] = train_acc1
train_acc5es[epoch] = train_acc5
train_times[epoch] = train_tm
lrs[epoch] = lr
with torch.no_grad():
for key, xloder in valid_loaders.items():
valid_loss, valid_acc1, valid_acc5, valid_tm = procedure(
xloder, network, criterion, None, None, "valid")
valid_losses["{:}@{:}".format(key, epoch)] = valid_loss
valid_acc1es["{:}@{:}".format(key, epoch)] = valid_acc1
valid_acc5es["{:}@{:}".format(key, epoch)] = valid_acc5
valid_times["{:}@{:}".format(key, epoch)] = valid_tm
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
need_time = "Time Left: {:}".format(
convert_secs2time(epoch_time.avg * (total_epoch - epoch - 1),
True))
logger.log(
"{:} {:} epoch={:03d}/{:03d} :: Train [loss={:.5f}, acc@1={:.2f}%, acc@5={:.2f}%] Valid [loss={:.5f}, acc@1={:.2f}%, acc@5={:.2f}%], lr={:}"
.format(
time_string(),
need_time,
epoch,
total_epoch,
train_loss,
train_acc1,
train_acc5,
valid_loss,
valid_acc1,
valid_acc5,
lr,
))
info_seed = {
"flop": flop,
"param": param,
"arch_config": arch_config._asdict(),
"opt_config": opt_config._asdict(),
"total_epoch": total_epoch,
"train_losses": train_losses,
"train_acc1es": train_acc1es,
"train_acc5es": train_acc5es,
"train_times": train_times,
"valid_losses": valid_losses,
"valid_acc1es": valid_acc1es,
"valid_acc5es": valid_acc5es,
"valid_times": valid_times,
"learning_rates": lrs,
"net_state_dict": net.state_dict(),
"net_string": "{:}".format(net),
"finish-train": True,
}
return info_seed
def main(xargs):
assert torch.cuda.is_available(), "CUDA is not available."
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
torch.set_num_threads(xargs.workers)
prepare_seed(xargs.rand_seed)
logger = prepare_logger(args)
train_data, valid_data, xshape, class_num = get_datasets(
xargs.dataset, xargs.data_path, -1
)
config = load_config(
xargs.config_path, {"class_num": class_num, "xshape": xshape}, logger
)
search_loader, _, valid_loader = get_nas_search_loaders(
train_data,
valid_data,
xargs.dataset,
"configs/nas-benchmark/",
(config.batch_size, config.test_batch_size),
xargs.workers,
)
logger.log(
"||||||| {:10s} ||||||| Search-Loader-Num={:}, Valid-Loader-Num={:}, batch size={:}".format(
xargs.dataset, len(search_loader), len(valid_loader), config.batch_size
)
)
logger.log("||||||| {:10s} ||||||| Config={:}".format(xargs.dataset, config))
search_space = get_search_spaces("cell", xargs.search_space_name)
model_config = dict2config(
{
"name": "RANDOM",
"C": xargs.channel,
"N": xargs.num_cells,
"max_nodes": xargs.max_nodes,
"num_classes": class_num,
"space": search_space,
"affine": False,
"track_running_stats": bool(xargs.track_running_stats),
},
None,
)
search_model = get_cell_based_tiny_net(model_config)
w_optimizer, w_scheduler, criterion = get_optim_scheduler(
search_model.parameters(), config
)
logger.log("w-optimizer : {:}".format(w_optimizer))
logger.log("w-scheduler : {:}".format(w_scheduler))
logger.log("criterion : {:}".format(criterion))
if xargs.arch_nas_dataset is None:
api = None
else:
api = API(xargs.arch_nas_dataset)
logger.log("{:} create API = {:} done".format(time_string(), api))
last_info, model_base_path, model_best_path = (
logger.path("info"),
logger.path("model"),
logger.path("best"),
)
network, criterion = torch.nn.DataParallel(search_model).cuda(), criterion.cuda()
if last_info.exists(): # automatically resume from previous checkpoint
logger.log(
"=> loading checkpoint of the last-info '{:}' start".format(last_info)
)
last_info = torch.load(last_info)
start_epoch = last_info["epoch"]
checkpoint = torch.load(last_info["last_checkpoint"])
genotypes = checkpoint["genotypes"]
valid_accuracies = checkpoint["valid_accuracies"]
search_model.load_state_dict(checkpoint["search_model"])
w_scheduler.load_state_dict(checkpoint["w_scheduler"])
w_optimizer.load_state_dict(checkpoint["w_optimizer"])
logger.log(
"=> loading checkpoint of the last-info '{:}' start with {:}-th epoch.".format(
last_info, start_epoch
)
)
else:
logger.log("=> do not find the last-info file : {:}".format(last_info))
start_epoch, valid_accuracies, genotypes = 0, {"best": -1}, {}
# start training
start_time, search_time, epoch_time, total_epoch = (
time.time(),
AverageMeter(),
AverageMeter(),
config.epochs + config.warmup,
)
for epoch in range(start_epoch, total_epoch):
w_scheduler.update(epoch, 0.0)
need_time = "Time Left: {:}".format(
convert_secs2time(epoch_time.val * (total_epoch - epoch), True)
)
epoch_str = "{:03d}-{:03d}".format(epoch, total_epoch)
logger.log(
"\n[Search the {:}-th epoch] {:}, LR={:}".format(
epoch_str, need_time, min(w_scheduler.get_lr())
)
)
# selected_arch = search_find_best(valid_loader, network, criterion, xargs.select_num)
search_w_loss, search_w_top1, search_w_top5 = search_func(
search_loader,
network,
criterion,
w_scheduler,
w_optimizer,
epoch_str,
xargs.print_freq,
logger,
)
search_time.update(time.time() - start_time)
logger.log(
"[{:}] searching : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%, time-cost={:.1f} s".format(
epoch_str, search_w_loss, search_w_top1, search_w_top5, search_time.sum
)
)
valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
valid_loader, network, criterion
)
logger.log(
"[{:}] evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%".format(
epoch_str, valid_a_loss, valid_a_top1, valid_a_top5
)
)
cur_arch, cur_valid_acc = search_find_best(
valid_loader, network, xargs.select_num
)
logger.log(
"[{:}] find-the-best : {:}, accuracy@1={:.2f}%".format(
epoch_str, cur_arch, cur_valid_acc
)
)
genotypes[epoch] = cur_arch
# check the best accuracy
valid_accuracies[epoch] = valid_a_top1
if valid_a_top1 > valid_accuracies["best"]:
valid_accuracies["best"] = valid_a_top1
find_best = True
else:
find_best = False
# save checkpoint
save_path = save_checkpoint(
{
"epoch": epoch + 1,
"args": deepcopy(xargs),
"search_model": search_model.state_dict(),
"w_optimizer": w_optimizer.state_dict(),
"w_scheduler": w_scheduler.state_dict(),
"genotypes": genotypes,
"valid_accuracies": valid_accuracies,
},
model_base_path,
logger,
)
last_info = save_checkpoint(
{
"epoch": epoch + 1,
"args": deepcopy(args),
"last_checkpoint": save_path,
},
logger.path("info"),
logger,
)
if find_best:
logger.log(
"<<<--->>> The {:}-th epoch : find the highest validation accuracy : {:.2f}%.".format(
epoch_str, valid_a_top1
)
)
copy_checkpoint(model_base_path, model_best_path, logger)
if api is not None:
logger.log("{:}".format(api.query_by_arch(genotypes[epoch], "200")))
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log("\n" + "-" * 200)
logger.log("Pre-searching costs {:.1f} s".format(search_time.sum))
start_time = time.time()
best_arch, best_acc = search_find_best(valid_loader, network, xargs.select_num)
search_time.update(time.time() - start_time)
logger.log(
"RANDOM-NAS finds the best one : {:} with accuracy={:.2f}%, with {:.1f} s.".format(
best_arch, best_acc, search_time.sum
)
)
if api is not None:
logger.log("{:}".format(api.query_by_arch(best_arch, "200")))
logger.close()
def simplify(save_dir, meta_file, basestr, target_dir):
meta_infos = torch.load(meta_file, map_location="cpu")
meta_archs = meta_infos["archs"] # a list of architecture strings
meta_num_archs = meta_infos["total"]
meta_max_node = meta_infos["max_node"]
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"))
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)."
.format(time_string(), num_evaluated_arch, meta_num_archs))
for key in sorted(list(num_seeds.keys())):
print(
"{:} There are {:5d} architectures that are evaluated {:} times.".
format(time_string(), num_seeds[key], key))
dataloader_dict = GET_DataLoaders(6)
to_save_simply = save_dir / "simplifies"
to_save_allarc = save_dir / "simplifies" / "architectures"
if not to_save_simply.exists():
to_save_simply.mkdir(parents=True, exist_ok=True)
if not to_save_allarc.exists():
to_save_allarc.mkdir(parents=True, exist_ok=True)
assert (save_dir /
target_dir) in subdir2archs, "can not find {:}".format(target_dir)
arch2infos, datasets = {}, (
"cifar10-valid",
"cifar10",
"cifar100",
"ImageNet16-120",
)
evaluated_indexes = set()
target_directory = save_dir / target_dir
target_less_dir = save_dir / "{:}-LESS".format(target_dir)
arch_indexes = subdir2archs[target_directory]
num_seeds = defaultdict(lambda: 0)
end_time = time.time()
arch_time = AverageMeter()
for idx, arch_index in enumerate(arch_indexes):
checkpoints = list(
target_directory.glob("arch-{:}-seed-*.pth".format(arch_index)))
ckps_less = list(
target_less_dir.glob("arch-{:}-seed-*.pth".format(arch_index)))
# create the arch info for each architecture
try:
arch_info_full = account_one_arch(
arch_index,
meta_archs[int(arch_index)],
checkpoints,
datasets,
dataloader_dict,
)
arch_info_less = account_one_arch(
arch_index,
meta_archs[int(arch_index)],
ckps_less,
["cifar10-valid"],
dataloader_dict,
)
num_seeds[len(checkpoints)] += 1
except:
print("Loading {:} failed, : {:}".format(arch_index, checkpoints))
continue
assert (int(arch_index) not in evaluated_indexes
), "conflict arch-index : {:}".format(arch_index)
assert (0 <= int(arch_index) < len(meta_archs)
), "invalid arch-index {:} (not found in meta_archs)".format(
arch_index)
arch_info = {"full": arch_info_full, "less": arch_info_less}
evaluated_indexes.add(int(arch_index))
arch2infos[int(arch_index)] = arch_info
torch.save(
{
"full": arch_info_full.state_dict(),
"less": arch_info_less.state_dict()
},
to_save_allarc / "{:}-FULL.pth".format(arch_index),
)
arch_info["full"].clear_params()
arch_info["less"].clear_params()
torch.save(
{
"full": arch_info_full.state_dict(),
"less": arch_info_less.state_dict()
},
to_save_allarc / "{:}-SIMPLE.pth".format(arch_index),
)
# measure elapsed time
arch_time.update(time.time() - end_time)
end_time = time.time()
need_time = "{:}".format(
convert_secs2time(arch_time.avg * (len(arch_indexes) - idx - 1),
True))
print("{:} {:} [{:03d}/{:03d}] : {:} still need {:}".format(
time_string(), target_dir, idx, len(arch_indexes), arch_index,
need_time))
# measure time
xstrs = [
"{:}:{:03d}".format(key, num_seeds[key])
for key in sorted(list(num_seeds.keys()))
]
print("{:} {:} done : {:}".format(time_string(), target_dir, xstrs))
final_infos = {
"meta_archs": meta_archs,
"total_archs": meta_num_archs,
"basestr": basestr,
"arch2infos": arch2infos,
"evaluated_indexes": evaluated_indexes,
}
save_file_name = to_save_simply / "{:}.pth".format(target_dir)
torch.save(final_infos, save_file_name)
print("Save {:} / {:} architecture results into {:}.".format(
len(evaluated_indexes), meta_num_archs, save_file_name))
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(
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(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 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 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(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 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(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)
valid_use = False
user_data = np.load(
'../../exps/NAS-Bench-201-algos/Dirichlet_100000000_Use_valid_{}_{}_non_iid_setting.npy'
.format(valid_use, args.dataset),
allow_pickle=True).item()
train_loader_list = {}
valid_loader_list = {}
# alignment_loader = torch.utils.data.DataLoader(
# DatasetSplit(train_data, np.random.choice(list(range(len(train_data))), 5000)),
# batch_size=args.batch_size,
# shuffle=True,
# num_workers=args.workers,
# pin_memory=True,
# )
alignment_loader = torch.utils.data.DataLoader(
DatasetSplit(train_data, user_data['public']),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
)
user_num = len(user_data) - 1
for user in range(user_num):
train_loader_list[user] = torch.utils.data.DataLoader(
DatasetSplit(train_data,
user_data[user]['train'] + user_data[user]['test']),
batch_size=args.batch_size,
shuffle=True,
drop_last=True,
num_workers=args.workers,
pin_memory=True,
)
valid_loader_list[user] = torch.utils.data.DataLoader(
DatasetSplit(valid_data, user_data[user]['valid']),
batch_size=args.batch_size,
shuffle=True,
drop_last=True,
num_workers=args.workers,
pin_memory=True,
)
# train_loader = torch.utils.data.DataLoader(
# train_data,
# batch_size=args.batch_size,
# shuffle=True,
# num_workers=args.workers,
# pin_memory=True,
# )
# valid_loader = torch.utils.data.DataLoader(
# valid_data,
# batch_size=args.batch_size,
# shuffle=False,
# num_workers=args.workers,
# pin_memory=True,
# )
# get configures
model_config = load_config(args.model_config, {"class_num": class_num},
logger)
optim_config = load_config(args.optim_config, {"class_num": class_num},
logger)
if args.model_source == "normal":
base_model = obtain_model(model_config)
elif args.model_source == "nas":
base_model = obtain_nas_infer_model(model_config,
args.extra_model_path)
elif args.model_source == "autodl-searched":
import ast
import re
file_proposal = args.extra_model_path
genotype_list = {}
if args.extra_model_path in Networks:
for user in range(user_num):
genotype_list[user] = Networks[args.extra_model_path]
else:
user_list = {}
user = 0
for line in open(file_proposal):
if "<<<--->>>" in line:
tep_dict = ast.literal_eval(
re.search('({.+})', line).group(0))
count = 0
for j in tep_dict['normal']:
for k in j:
if 'skip_connect' in k[0]:
count += 1
if count == 2:
# if user%5 not in genotype_list:
# logger.log("user{}'s architecture is chosen from epoch {}".format(user%5, user//5))
genotype_list[user % 5] = tep_dict
user_list[user % 5] = user // 5
user += 1
for user in user_list:
logger.log(
"user{}'s architecture is chosen from epoch {}".format(
user, user_list[user]))
logger.log(genotype_list)
base_model_list = {}
for user in range(user_num):
base_model_list[user] = obtain_model(model_config,
genotype_list[3])
flop, param = get_model_infos(base_model_list[user], xshape)
logger.log("The model of User {}: parm: {}, Flops: {}.".format(
user, param, flop))
wandb.watch(base_model_list[user])
# base_model = obtain_model(model_config, args.extra_model_path)
elif args.model_source == "Densenet":
base_model_list = {}
for user in range(user_num):
base_model_list[user] = torch.hub.load('pytorch/vision:v0.10.0',
'densenet121',
pretrained=False)
flop, param = get_model_infos(base_model_list[user], xshape)
logger.log("The model of User {}: parm: {}, Flops: {}.".format(
user, param, flop))
else:
base_model_list = {}
for user in range(user_num):
base_model_list[user], _, __ = create_cnn_model(
args.model_source,
args.dataset,
optim_config.epochs + optim_config.warmup,
None,
use_cuda=1)
flop, param = get_model_infos(base_model_list[user], xshape)
logger.log("The model of User {}: parm: {}, Flops: {}.".format(
user, param, flop))
# raise ValueError("invalid model-source : {:}".format(args.model_source))
optimizer_list = {}
scheduler_list = {}
criterion_list = {}
for user in range(user_num):
flop, param = get_model_infos(base_model_list[user], xshape)
# logger.log("model ====>>>>:\n{:}".format(base_model_list[user]))
# logger.log("model information : {:}".format(base_model_list[user].get_message()))
logger.log("-" * 50)
logger.log("Params={:.2f} MB, FLOPs={:.2f} M ... = {:.2f} G".format(
param, flop, flop / 1e3))
logger.log("-" * 50)
optimizer_list[user], scheduler_list[user], criterion_list[
user] = get_optim_scheduler(base_model_list[user].parameters(),
optim_config)
# logger.log("User{}, train_data : {:}".format(user, train_data[user]))
# logger.log("User{}, valid_data : {:}".format(user, valid_data[user]))
# optimizer, scheduler, criterion = get_optim_scheduler(
# base_model.parameters(), optim_config
# )
logger.log("User{}, optimizer : {:}".format(user,
optimizer_list[user]))
logger.log("User{}, scheduler : {:}".format(user,
scheduler_list[user]))
logger.log("User{}, criterion : {:}".format(user,
criterion_list[user]))
# base_model_list[user], criterion_list[user] = torch.nn.DataParallel(base_model[user]).cuda(), criterion_list[user].cuda()
criterion_list[user] = criterion_list[user].cuda()
base_model_list[user] = base_model_list[user].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_infox = torch.load(last_info)
start_epoch = last_infox["epoch"] + 1
last_checkpoint_path = last_infox["last_checkpoint"]
if not last_checkpoint_path.exists():
logger.log("Does not find {:}, try another path".format(
last_checkpoint_path))
last_checkpoint_path = (last_info.parent /
last_checkpoint_path.parent.name /
last_checkpoint_path.name)
checkpoint = torch.load(last_checkpoint_path)
for user in base_model_list:
base_model_list[user].load_state_dict(
checkpoint["model_{}".format(user)])
optimizer_list[user].load_state_dict(
checkpoint["optimizer_{}".format(user)])
scheduler_list[user].load_state_dict(
checkpoint["scheduler_{}".format(user)])
valid_accuracies = checkpoint["valid_accuracies"]
logger.log(
"=> loading checkpoint of the last-info '{:}' start with {:}-th epoch."
.format(last_info, start_epoch))
del (checkpoint)
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
for user in base_model_list:
base_model_list[user].load_state_dict(
checkpoint["model_{}".format(user)])
optimizer_list[user].load_state_dict(
checkpoint["optimizer_{}".format(user)])
scheduler_list[user].load_state_dict(
checkpoint["scheduler_{}".format(user)])
valid_accuracies = checkpoint["valid_accuracies"]
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
local_epoch = args.local_epoch
# Main Training and Evaluation Loop
start_time = time.time()
epoch_time = AverageMeter()
for epoch in range(start_epoch, total_epoch):
epoch_str = "epoch={:03d}/{:03d}".format(epoch, total_epoch)
test_accuracy1_list = []
test_accuracy5_list = []
for user in scheduler_list:
if (epoch % 1 == 0) or (epoch + 1 == total_epoch):
logger.log("-" * 150)
valid_loss, valid_acc1, valid_acc5 = valid_func(
valid_loader_list[user],
base_model_list[user],
criterion_list[user],
optim_config,
epoch_str,
args.print_freq_eval,
logger,
)
logger.log(
"Important: User {}: ***{:s}*** VALID [{:}] loss = {:.6f}, accuracy@1 = {:.2f}, accuracy@5 = {:.2f} | Best-Valid-Acc@1={:.2f}, Error@1={:.2f}"
.format(
user,
time_string(),
epoch_str,
valid_loss,
valid_acc1,
valid_acc5,
valid_accuracies["best"],
100 - valid_accuracies["best"],
))
test_accuracy1_list.append(valid_acc1)
test_accuracy5_list.append(valid_acc5)
if args.logits_aggregation:
Logits_aggregation_func(alignment_loader, base_model_list,
optimizer_list, logger, 3)
else:
tep_list = [
model.state_dict() for model in base_model_list.values()
]
global_state = average_weights(tep_list)
del (tep_list)
for one in base_model_list:
base_model_list[one].load_state_dict(global_state)
for user in scheduler_list:
scheduler_list[user].update(epoch, 0.0)
need_time = "Time Left: {:}".format(
convert_secs2time(epoch_time.avg * (total_epoch - epoch), True))
LRs = scheduler_list[0].get_lr()
find_best = False
# set-up drop-out ratio
# if hasattr(base_model, "update_drop_path"):
# base_model.update_drop_path(
# model_config.drop_path_prob * epoch / total_epoch
# )
logger.log(
"\n***{:s}*** start {:s} {:s}, LR=[{:.12f} ~ {:.12f}], scheduler={:}"
.format(time_string(), epoch_str, need_time, min(LRs), max(LRs),
scheduler_list[0]))
# train for one epoch
for user in train_loader_list:
train_loss, train_acc1, train_acc5 = train_func(
train_loader_list[user], base_model_list[user],
criterion_list[user], scheduler_list[user],
optimizer_list[user], optim_config, epoch_str, args.print_freq,
logger, local_epoch)
# log the results
logger.log(
"User {} ***{:s}*** TRAIN [{:}] loss = {:.6f}, accuracy-1 = {:.2f}, accuracy-5 = {:.2f}"
.format(user, time_string(), epoch_str, train_loss, train_acc1,
train_acc5))
info_dict = {
"{}user_train_loss".format(user): train_loss,
"{}user_train_top1".format(user): train_acc1,
"{}user_train_top5".format(user): train_acc5,
"{}user_valid_loss".format(user): valid_loss,
"{}user_valid_top1".format(user): valid_acc1,
"{}user_valid_top5".format(user): valid_acc5,
"epoch": epoch
}
wandb.log(info_dict)
if np.average(test_accuracy1_list) > valid_accuracies["best"]:
valid_accuracies["best"] = np.average(test_accuracy1_list)
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,
))
valid_accuracies[epoch] = np.average(test_accuracy1_list)
info_dict = {
"average_valid_top1_acc": np.average(test_accuracy1_list),
"average_valid_top5_acc": np.average(test_accuracy5_list),
"epoch": epoch
}
wandb.log(info_dict)
# 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
checkpoint_dict = {
"epoch": epoch,
"args": deepcopy(args),
"FLOP": flop,
"PARAM": param,
"model_source": args.model_source,
"valid_accuracies": deepcopy(valid_accuracies),
"model-config": model_config._asdict(),
"optim-config": optim_config._asdict()
}
for user in base_model_list:
checkpoint_dict["model_{}".format(
user)] = base_model_list[user].state_dict()
checkpoint_dict["scheduler_{}".format(
user)] = scheduler_list[user].state_dict()
checkpoint_dict["optimizer_{}".format(
user)] = optimizer_list[user].state_dict()
save_path = save_checkpoint(checkpoint_dict, model_base_path, logger)
del (checkpoint_dict)
if find_best:
copy_checkpoint(model_base_path, model_best_path, logger)
last_info = save_checkpoint(
{
"epoch": epoch,
"args": deepcopy(args),
"last_checkpoint": save_path,
},
logger.path("info"),
logger,
)
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
logger.log("\n" + "-" * 200)
# logger.log(
# "Finish training/validation in {:} with Max-GPU-Memory of {:.2f} MB, and save final checkpoint into {:}".format(
# convert_secs2time(epoch_time.sum, True),
# max(v for k, v in max_bytes.items()) / 1e6,
# logger.path("info"),
# )
# )
logger.log("-" * 200 + "\n")
logger.close()
