def _get_arch_param(self,
generator,
hardware_constraint=None,
valid=False):
# ====================== Strict fair sample
if hardware_constraint is None:
hardware_constraint = torch.tensor(
self.hardware_pool[self.hardware_index] + random.random() -
0.5,
dtype=torch.float32).view(-1, 1)
#hardware_constraint = torch.tensor(self.hardware_pool[self.hardware_index], dtype=torch.float32).view(-1, 1)
self.hardware_index += 1
if self.hardware_index == len(self.hardware_pool):
self.hardware_index = 0
random.shuffle(self.hardware_pool)
else:
hardware_constraint = torch.tensor(hardware_constraint,
dtype=torch.float32).view(
-1, 1)
# ======================
hardware_constraint = hardware_constraint.to(self.device)
logging.info("Target macs : {}".format(hardware_constraint.item()))
normalize_hardware_constraint = min_max_normalize(
self.CONFIG.high_macs, self.CONFIG.low_macs, hardware_constraint)
noise = torch.randn(*self.backbone.shape)
noise = noise.to(self.device)
noise *= self.noise_weight
arch_param = generator(self.backbone, normalize_hardware_constraint,
noise)
return hardware_constraint, arch_param
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 _get_arch_param(self, generator, target_hardware_constraint=None):
"""
Given the target hardware constraint as the input of generator. The generator output the architecture parameter.
"""
hardware_constraint = target_hardware_constraint.to(self.device)
logging.info("Target flops : {}".format(hardware_constraint.item()))
prior = self.prior_pool.get_prior(hardware_constraint.item())
prior = prior.to(self.device)
normalize_hardware_constraint = min_max_normalize(
self.CONFIG.high_flops, self.CONFIG.low_flops, hardware_constraint)
arch_param = generator(prior, normalize_hardware_constraint)
return arch_param
def evaluate_generator(generator,
prior_pool,
lookup_table,
CONFIG,
device,
val=True):
"""
Evaluate kendetall and hardware constraint loss of generator
"""
total_loss = 0
evaluate_metric = {"gen_flops": [], "true_flops": []}
for mac in range(CONFIG.low_flops, CONFIG.high_flops, 10):
hardware_constraint = torch.tensor(mac, dtype=torch.float32)
hardware_constraint = hardware_constraint.view(-1, 1)
hardware_constraint = hardware_constraint.to(device)
prior = prior_pool.get_prior(hardware_constraint.item())
prior = prior.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)
layers_config = lookup_table.decode_arch_param(arch_param)
gen_mac = lookup_table.get_model_flops(arch_param)
hc_loss = cal_hc_loss(gen_mac.cuda(), hardware_constraint.item(),
CONFIG.alpha, CONFIG.loss_penalty)
evaluate_metric["gen_flops"].append(gen_mac.item())
evaluate_metric["true_flops"].append(mac)
total_loss += hc_loss.item()
tau, _ = stats.kendalltau(evaluate_metric["gen_flops"],
evaluate_metric["true_flops"])
return evaluate_metric, total_loss, tau
def search_train_loop(self, generator):
self.epochs = self.warmup_epochs + self.search_epochs
# Training generator
best_loss = 10000.0
best_top1 = 0
tau = 5
for epoch in range(self.warmup_epochs, self.search_epochs):
logging.info("Start to train for search epoch {}".format(epoch))
logging.info("Tau: {}".format(tau))
self._generator_training_step(generator,
val_loader,
epoch,
tau,
info_for_logger="_gen_train_step")
# ================ Train ============================================
for i in range():
# Training generator
arch_param, hardware_constraint = self.set_arch_param(
generator, tau=tau)
# ============== evaluation flops ===============================
gen_flops = self.flops_table.predict_arch_param_efficiency(
arch_param)
hc_loss = cal_hc_loss(gen_flops.cuda(),
hardware_constraint.item(),
self.CONFIG.alpha,
self.CONFIG.loss_penalty)
# ===============================================================
self.g_optimizer.zero_grad()
# ============== predict top1 accuracy ==========================
top1_avg = self.accuracy_predictor(arch_param)
ce_loss = -1 * top1_avg
# ===============================================================
loss = ce_loss + hc_loss
logging.info("HC loss : {}".format(hc_loss))
loss.backward()
self.g_optimizer.step()
self.g_optimizer.zero_grad()
# ====================================================================
# ============== Valid ===============================================
hardware_constraint, arch_param = self._get_arch_param(
generator, hardware_constraint, valid=True)
arch_param = self.calculate_one_hot(arch_param)
arch_param, hardware_constraint = self.set_arch_param(
generator,
model,
hardware_constraint=hardware_constraint,
arch_param=arch_param)
# ============== evaluation flops ===============================
gen_flops = self.flops_table.predict_arch_param_efficiency(
arch_param)
hc_loss = cal_hc_loss(gen_flops.cuda(), hardware_constraint.item(),
self.CONFIG.alpha, self.CONFIG.loss_penalty)
# ===============================================================
# ============== predict top1 accuracy ==========================
top1_avg = self.accuracy_predictor(arch_param)
logger.info("Valid : Top-1 avg : {}".format(top1_avg))
# ===============================================================
# ====================================================================
# ============== Evaluate ============================================
total_loss = 0
evaluate_metric = {"gen_flops": [], "true_flops": []}
for flops in range(self.CONFIG.low_macs, self.CONFIG.high_macs,
10):
hardware_constraint = torch.tensor(flops, dtpye=torch.float32)
hardware_constraint = hardware_constraint.view(-1, 1)
hardware_constraint = hardware_constraint.to(self.device)
normalize_hardware_constraint = min_max_normalize(
self.CONFIG.high_macs, self.CONFIG.low_macs,
hardware_constraint)
noise = torch.randn(*self.backbone.shape)
noise = noise.to(device)
noise *= 0
arch_param = generator(self.backbone,
normalize_hardware_constraint, noise)
# ============== evaluation flops ===============================
gen_flops = self.flops_table.predict_arch_param_efficiency(
arch_param)
hc_loss = cal_hc_loss(gen_flops.cuda(),
hardware_constraint.item(),
self.CONFIG.alpha,
self.CONFIG.loss_penalty)
# ===============================================================
evaluate_metric["gen_flops"].append(gen_flops)
evaluate_metric["true_flops"].append(flops)
total_loss += hc_loss.item()
kendall_tau, _ = stats.kendalltau(evaluate_metric["gen_flops"],
evaluate_metric["true_flops"])
# ====================================================================
logging.info("Total loss : {}".format(total_loss))
if best_loss > total_loss:
logging.info("Best loss by now: {} Tau : {}.Save model".format(
total_loss, kendall_tau))
best_loss = total_loss
save_generator_evaluate_metric(
evaluate_metric, self.CONFIG.path_to_generator_eval)
save(generator, self.g_optimizer,
self.CONFIG.path_to_save_generator)
if top1_avg > best_top1 and total_loss < 0.4:
logging.info(
"Best top1-avg by now: {}.Save model".format(top1_avg))
best_top1 = top1_avg
save(generator, self.g_optimizer,
self.CONFIG.path_to_best_avg_generator)
save(generator, self.g_optimizer,
"./logs/generator/{}.pth".format(total_loss))
tau *= self.CONFIG.tau_decay
self.noise_weight = self.noise_weight * self.CONFIG.noise_decay if self.noise_weight > 0.0001 else 0
logging.info("Noise weight : {}".format(self.noise_weight))
logging.info("Best loss: {}".format(best_loss))
save(generator, self.g_optimizer, self.CONFIG.path_to_fianl_generator)
generator.to(device)
prior_pool = PriorPool(lookup_table, arch_param_nums, None, None, None,
CONFIG)
# Sample architecture parameter =======================
prior = prior_pool.get_prior(args.flops)
prior = prior.to(device)
noise = torch.randn(*prior.shape)
noise = noise.to(device)
noise *= 0
hardware_constraint = torch.tensor(args.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)
cal_model_efficient(model, CONFIG)
if (device.type == "cuda" and CONFIG.ngpu >= 1):
