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

Code before:
    parser.add_argument("--log-frequency", type=int, default=10)
    parser.add_argument("--label-smoothing", type=float, default=0.1)
    parser.add_argument("--evolution-sample-size", type=int, default=10)
    parser.add_argument("--evolution-population-size", type=int, default=50)
    parser.add_argument("--evolution-cycles", type=int, default=10)

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


Review comment (suggestion):
would also be better to add comment to tell users that why `nni.trace` is used here. to tell users that `nni.trace` is used to make these object serializable, so that the trials can run other other processes or on remote servers.

Code after:
    parser.add_argument("--test-batch-size", type=int, default=512)
    parser.add_argument("--log-frequency", type=int, default=10)
    parser.add_argument("--label-smoothing", type=float, default=0.1)
    parser.add_argument("--evolution-sample-size", type=int, default=10)
    parser.add_argument("--evolution-population-size", type=int, default=50)
    parser.add_argument("--evolution-cycles", type=int, default=10)
    parser.add_argument("--latency-filter", type=str, default=None,
                        help="Apply latency filter by calling the name of the applied hardware.")
    parser.add_argument("--latency-threshold", type=float, default=100)

    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)