Repository: microsoft/nni
File: examples/nas/oneshot/spos/search.py (Python)

Code before:
    torch.backends.cudnn.deterministic = True

    assert torch.cuda.is_available()

    base_model = ShuffleNetV2OneShot()
    base_latency_predictor = 'cortexA76cpu_tflite21'
    criterion = CrossEntropyLabelSmooth(1000, args.label_smoothing)
    if args.spos_preprocessing:
        trans = nni.trace(transforms.Compose)([
            nni.trace(transforms.RandomResizedCrop)(224),
            nni.trace(transforms.ColorJitter)(brightness=0.4, contrast=0.4, saturation=0.4),
            nni.trace(transforms.RandomHorizontalFlip)(0.5),
            nni.trace(ToBGRTensor)(),
        ])
    else:
        trans = nni.trace(transforms.Compose)([
            nni.trace(transforms.RandomResizedCrop)(224),
            nni.trace(transforms.ToTensor)()
        ])
    train_dataset = nni.trace(datasets.ImageNet)(args.imagenet_dir, split='train', transform=trans)
    val_dataset = nni.trace(datasets.ImageNet)(args.imagenet_dir, split='val', transform=trans)

    evaluator = FunctionalEvaluator(evaluate_acc, criterion=criterion, args=args, train_dataset=train_dataset, val_dataset=val_dataset)
    evolution_strategy = strategy.RegularizedEvolution(
        # model_filter=LatencyFilter(threshold=100, predictor=base_latency_predictor),
        sample_size=1, population_size=2, cycles=2)
    exp = RetiariiExperiment(base_model, evaluator, strategy=evolution_strategy)

    exp_config = RetiariiExeConfig('local')
    exp_config.trial_concurrency = 2
    exp_config.trial_gpu_number = 1
    exp_config.training_service.use_active_gpu = False
    exp_config.execution_engine = 'base'
    # exp_config.dummy_input = [1, 3, 224, 224]

    exp.run(exp_config, 8084)

    print('Exported models:')
    for i, model in enumerate(exp.export_top_models(formatter='dict')):
        print(model)
        with open(f'architecture_final_{i}.json', 'w') as f: 
            json.dump(get_archchoice_by_model(model), f, indent=4)

if __name__ == "__main__":
    _main()


Review comment (suggestion):
You can make these hyper-parameters appear in argparse.

By default, I think `sample_size=10`, `population_size=50`, `cycles=10`.

Code after:

    args = parser.parse_args()

    # use a fixed set of image will improve the performance
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    np.random.seed(args.seed)
    random.seed(args.seed)
    torch.backends.cudnn.deterministic = True

    assert torch.cuda.is_available()

    base_model = ShuffleNetV2OneShot()
    criterion = CrossEntropyLabelSmooth(1000, args.label_smoothing)
    if args.spos_preprocessing:
        # ``nni.trace`` is used to make transforms serializable, so that the trials can run other processes or on remote servers.
        trans = nni.trace(transforms.Compose)([
            nni.trace(transforms.RandomResizedCrop)(224),
            nni.trace(transforms.ColorJitter)(brightness=0.4, contrast=0.4, saturation=0.4),
            nni.trace(transforms.RandomHorizontalFlip)(0.5),
            nni.trace(ToBGRTensor)(),
        ])
    else:
        # ``nni.trace`` is used to make transforms serializable, so that the trials can run other processes or on remote servers.
        trans = nni.trace(transforms.Compose)([
            nni.trace(transforms.RandomResizedCrop)(224),
            nni.trace(transforms.ToTensor)()
        ])
    train_dataset = nni.trace(datasets.ImageNet)(args.imagenet_dir, split='train', transform=trans)
    val_dataset = nni.trace(datasets.ImageNet)(args.imagenet_dir, split='val', transform=trans)

    if args.latency_filter:
        latency_filter = LatencyFilter(threshold=args.latency_threshold, predictor=args.latency_filter)
    else:
        latency_filter = None

    evaluator = FunctionalEvaluator(evaluate_acc, criterion=criterion, args=args, train_dataset=train_dataset, val_dataset=val_dataset)
    evolution_strategy = strategy.RegularizedEvolution(
        model_filter=latency_filter,
        sample_size=args.evolution_sample_size, population_size=args.evolution_population_size, cycles=args.evolution_cycles)
    exp = RetiariiExperiment(base_model, evaluator, strategy=evolution_strategy)

    exp_config = RetiariiExeConfig('local')
    exp_config.trial_concurrency = 2
    exp_config.trial_gpu_number = 1
    exp_config.max_trial_number = args.evolution_cycles
    exp_config.training_service.use_active_gpu = False
    exp_config.execution_engine = 'base'
    exp_config.dummy_input = [1, 3, 224, 224]

    exp.run(exp_config, args.port)