One-shot NAS

Before reading this tutorial, we highly recommend you to first go through the tutorial of how to define a model space.

Model Search with One-shot Trainer

With a defined model space, users can explore the space in two ways. One is using strategy and single-arch evaluator as demonstrated here. The other is using one-shot trainer, which consumes much less computational resource compared to the first one. In this tutorial we focus on this one-shot approach. The principle of one-shot approach is combining all the models in a model space into one big model (usually called super-model or super-graph). It takes charge of both search, training and testing, by training and evaluating this big model.

We list the supported one-shot trainers here:

  • DARTS trainer

  • ENAS trainer

  • ProxylessNAS trainer

  • Single-path (random) trainer

See API reference for detailed usages. Here, we show an example to use DARTS trainer manually.

from nni.retiarii.oneshot.pytorch import DartsTrainer
trainer = DartsTrainer(
    model=model,
    loss=criterion,
    metrics=lambda output, target: accuracy(output, target, topk=(1,)),
    optimizer=optim,
    num_epochs=args.epochs,
    dataset=dataset_train,
    batch_size=args.batch_size,
    log_frequency=args.log_frequency,
    unrolled=args.unrolled
)
trainer.fit()
final_architecture = trainer.export()

After the searching is done, we can use the exported architecture to instantiate the full network for retraining. Here is an example:

from nni.retiarii import fixed_arch
with fixed_arch('/path/to/checkpoint.json'):
    model = Model()