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))