コード例 #1
0
ファイル: model.py プロジェクト: eric8607242/SGNAS
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
コード例 #2
0
    if CONFIG.cuda:
        device = torch.device("cuda" if (torch.cuda.is_available() and CONFIG.ngpu > 0) else "cpu")
    else:
        device = torch.device("cpu")

    get_logger(CONFIG.log_dir)
    writer = get_writer(CONFIG.write_dir)

    #set_random_seed(CONFIG.seed)

    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)

    model = Supernet(CONFIG)
    lookup_table = LookUpTable(CONFIG)

    arch_param_nums = model.get_arch_param_nums()
    #generator = ConvGenerator(CONFIG.hc_dim, 1, CONFIG.hidden_dim)
    generator = get_generator(CONFIG, arch_param_nums)

    criterion = cross_encropy_with_label_smoothing

    if CONFIG.generator_pretrained is not None:
        logging.info("Loading model")
        model.load_state_dict(torch.load(CONFIG.model_pretrained)["model"])
        generator.load_state_dict(torch.load(CONFIG.generator_pretrained)["model"])

    generator.to(device)
    model.to(device)
コード例 #3
0
        device = torch.device("cuda" if (torch.cuda.is_available() and CONFIG.ngpu > 0) else "cpu")
    else:
        device = torch.device("cpu")

    get_logger(CONFIG.log_dir)
    writer = get_writer(args.title, CONFIG.write_dir)

    logging.info("=================================== Experiment title : {} Start ===========================".format(args.title))

    set_random_seed(CONFIG.seed)

    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)

    model = Supernet(CONFIG)
    lookup_table = LookUpTable(CONFIG)

    criterion = cross_encropy_with_label_smoothing

    arch_param_nums = model.get_arch_param_nums()               # layer_num * (sub_blocks * kernels_nums)
    generator = get_generator(CONFIG, arch_param_nums)

    # if CONFIG.model_pretrained is not None and os.path.isfile(CONFIG.model_pretrained):
    #     logging.info("Load pretrained weight from {}".format(CONFIG.model_pretrained))
    #     model.load_state_dict(torch.load(CONFIG.model_pretrained)["model"])

    model.to(device)
    generator.to(device)
    if (device.type == "cuda" and CONFIG.ngpu >= 1):
        model = nn.DataParallel(model, list(range(CONFIG.ngpu)))
コード例 #4
0
        nodes_num = calculate_nodes(CONFIG)
        adj_matrix = get_adj_matrix(nodes_num, CONFIG)
        adj_matrix = get_random_architecture(adj_matrix, CONFIG)

    if args.train_data:
        for i in range(100):
            total_top1 = 0
            for j in range(2):
                if args.load_architecture:
                    adj_matrix = architectures.iloc[i].values
                    adj_matrix = adj_matrix.reshape(nodes_num, nodes_num)
                else:
                    adj_matrix = get_adj_matrix(nodes_num, CONFIG)
                    adj_matrix = get_random_architecture(adj_matrix, CONFIG)

                model = Supernet(adj_matrix, CONFIG)
                model = model.to(device)
                if (device.type == "cuda" and CONFIG.ngpu >= 1):
                    model = nn.DataParallel(model, list(range(CONFIG.ngpu)))

                criterion = cross_encropy_with_label_smoothing
                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, None, device, CONFIG)
                best_top1 = trainer.train_loop(train_loader, test_loader, model)
                logging.info("Total training time : {:.2f}".format(time.time() - start_time))
コード例 #5
0
ファイル: train_cifar.py プロジェクト: eric8607242/SGNAS
    get_logger(CONFIG.log_dir)
    writer = get_writer(args.title, CONFIG.write_dir)

    logging.info(
        "=================================== Experiment title : {} Start ==========================="
        .format(args.title))

    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, test_dataset, test_dataset, CONFIG)

    lookup_table = LookUpTable(CONFIG)

    supernet = Supernet(CONFIG)
    arch_param_nums = supernet.get_arch_param_nums()

    generator = get_generator(CONFIG, arch_param_nums)

    criterion = cross_encropy_with_label_smoothing

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