def evaluate_lookup_table(lookup_table, prior_pool, CONFIG, evaluate_nums=10):
for i in range(evaluate_nums):
gen_mac, arch_param = prior_pool.generate_arch_param(lookup_table)
gen_mac = lookup_table.get_model_flops(arch_param.cuda())
layers_config = lookup_table.decode_arch_param(arch_param)
model = Model(layers_config, CONFIG.dataset, CONFIG.classes)
cal_model_efficient(model, CONFIG)
def get_model(config_path,
target_flops,
num_classes=1000,
in_chans=3,
activation="relu",
se=False,
bn_momentum=0.1):
CONFIG = get_config(config_path)
if CONFIG.cuda:
device = torch.device("cuda" if (
torch.cuda.is_available() and CONFIG.ngpu > 0) else "cpu")
else:
device = torch.device("cpu")
lookup_table = LookUpTable(CONFIG)
supernet = Supernet(CONFIG)
arch_param_nums = supernet.get_arch_param_nums()
generator = get_generator(CONFIG, arch_param_nums)
if CONFIG.generator_pretrained is not None:
generator.load_state_dict(
torch.load(CONFIG.generator_pretrained)["model"])
generator.to(device)
prior_pool = PriorPool(lookup_table, arch_param_nums, None, None, None,
CONFIG)
# Sample architecture parameter =======================
prior = prior_pool.get_prior(target_flops)
prior = prior.to(device)
hardware_constraint = torch.tensor(target_flops).to(device)
normalize_hardware_constraint = min_max_normalize(CONFIG.high_flops,
CONFIG.low_flops,
hardware_constraint)
arch_param = generator(prior, normalize_hardware_constraint)
arch_param = lookup_table.get_validation_arch_param(arch_param)
gen_flops = lookup_table.get_model_flops(arch_param)
logging.info("Generate flops : {}".format(gen_flops))
layers_config = lookup_table.decode_arch_param(arch_param)
model = Model(l_cfgs=layers_config,
dataset=CONFIG.dataset,
classes=CONFIG.classes,
activation=activation,
se=se,
bn_momentum=bn_momentum)
cal_model_efficient(model, CONFIG)
return model
def evaluate_arch_param(evaluate_trainer, supernet, generator, train_loader, test_loader, parameter_metric, backbone_pool, lookup_table, device, CONFIG):
avg_metric = {"gen_macs":[], "avg":[], "supernet_avg":[]}
for a in range(parameter_metric.shape[0]):
arch_param = torch.tensor(parameter_metric.iloc[a].values.reshape(19, -1), dtype=torch.float)
layers_config = lookup_table.decode_arch_param(arch_param)
model = Model(layers_config, CONFIG.dataset, CONFIG.classes)
model = model.to(device)
if (device.type == "cuda" and CONFIG.ngpu >= 1) :
model = nn.DataParallel(model, list(range(CONFIG.ngpu)))
macs = cal_model_efficient(model, CONFIG)
arch_param, hardware_constraint = evaluate_trainer.set_arch_param(generator, supernet, hardware_constraint=macs, arch_param=arch_param)
supernet_avg, _ = evaluate_trainer.generator_validate(generator, supernet, test_loader, 0, 0, hardware_constraint=hardware_constraint, arch_param=arch_param, sample=False)
avg_accuracy = 0
"""
for t in range(CONFIG.train_time):
model = Model(layers_config, CONFIG.dataset, CONFIG.classes)
model = model.to(device)
if (device.type == "cuda" and CONFIG.ngpu >= 1) :
model = nn.DataParallel(model, list(range(CONFIG.ngpu)))
optimizer = get_optimizer(model, CONFIG.optim_state)
scheduler = get_lr_scheduler(optimizer, len(train_loader), CONFIG)
trainer = Trainer(criterion, None, optimizer, None, scheduler, writer, device, lookup_table, backbone_pool, CONFIG)
top1_avg = trainer.train_loop(train_loader, test_loader, model)
avg_accuracy += top1_avg
"""
avg_accuracy /= CONFIG.train_time
avg_metric["gen_macs"].append(macs)
avg_metric["avg"].append(avg_accuracy)
avg_metric["supernet_avg"].append(supernet_avg)
return avg_metric
get_logger(CONFIG.log_dir)
writer = get_writer(CONFIG.write_dir)
train_transform, val_transform, test_transform = get_transforms(CONFIG)
train_dataset, val_dataset, test_dataset = get_dataset(
train_transform, val_transform, test_transform, CONFIG)
train_loader, val_loader, test_loader = get_dataloader(
train_dataset, val_dataset, test_dataset, CONFIG)
lookup_table = LookUpTable(CONFIG)
criterion = cross_encropy_with_label_smoothing
layers_config = lookup_table.decode_arch_param(arch_param)
model = Model(layers_config, CONFIG.dataset, CONFIG.classes)
model = model.to(device)
if (device.type == "cuda" and CONFIG.ngpu >= 1):
model = nn.DataParallel(model, list(range(CONFIG.ngpu)))
cal_model_efficient(model, CONFIG)
optimizer = get_optimizer(model, CONFIG.optim_state)
scheduler = get_lr_scheduler(optimizer, len(train_loader), CONFIG)
start_time = time.time()
trainer = Trainer(criterion, optimizer, scheduler, writer, device, CONFIG)
trainer.train_loop(train_loader, test_loader, model)
logging.info("Total training time : {:.2f}".format(time.time() -
start_time))