Exploration Strategy

There are two types of model space exploration approach: Multi-trial strategy and One-shot strategy. When the model space has been constructed, users can use either exploration approach to explore the model space.

Here is the list of exploration strategies that NNI has supported.



Brief Description



Randomly sample an architecture each time



Traverse the search space and try all possibilities



Evolution algorithm for NAS. Reference



Tree-structured Parzen Estimator (TPE). Reference



Policy-based reinforcement learning, based on implementation of tianshou. Reference



Continuous relaxation of the architecture representation, allowing efficient search of the architecture using gradient descent. Reference



RL controller learns to generate the best network on a super-net. Reference



Choose the best block by using Gumbel Softmax random sampling and differentiable training. Reference



Train a super-net with uniform path sampling. Reference



A low-memory-consuming optimized version of differentiable architecture search. Reference

Multi-trial strategy

Multi-trial NAS means each sampled model from model space is trained independently. A typical multi-trial NAS is NASNet. In multi-trial NAS, users need model evaluator to evaluate the performance of each sampled model, and need an exploration strategy to sample models from a defined model space. Here, users could use NNI provided model evaluators or write their own model evalutor. They can simply choose a exploration strategy. Advanced users can also customize new exploration strategy.

To use an exploration strategy, users simply instantiate an exploration strategy and pass the instantiated object to NasExperiment. Below is a simple example.

import nni.nas.strategy as strategy
exploration_strategy = strategy.Random(dedup=True)

Rather than using strategy.Random, users can choose one of the strategies from the table above.

One-shot strategy

One-shot NAS algorithms leverage weight sharing among models in neural architecture search space to train a supernet, and use this supernet to guide the selection of better models. This type of algorihtms greatly reduces computational resource compared to independently training each model from scratch (which we call “Multi-trial NAS”).

The usage of one-shot strategies are much alike to multi-trial strategies. Users simply need to create a strategy and run NasExperiment. Since one-shot strategies will manipulate the training recipe, to use a one-shot strategy, the evaluator needs to be one of the PyTorch-Lightning evaluators, either built-in or customized. Example follows:

import nni.nas.strategy as strategy
import nni.nas.evaluator.pytorch.lightning as pl
evaluator = pl.Classification(
  # Need to use `pl.DataLoader` instead of `torch.utils.data.DataLoader` here,
  # or use `nni.trace` to wrap `torch.utils.data.DataLoader`.
  train_dataloaders=pl.DataLoader(train_dataset, batch_size=100),
  val_dataloaders=pl.DataLoader(test_dataset, batch_size=100),
  # Other keyword arguments passed to pytorch_lightning.Trainer.
exploration_strategy = strategy.DARTS()

One-shot strategies only support a limited set of mutation primitives. See the reference for the detailed support list of each algorithm.

One-shot strategy is compatible with Lightning accelerators. It means that, you can accelerate one-shot strategies on hardwares like multiple GPUs. To enable this feature, you only need to pass the keyword arguments which used to be set in pytorch_lightning.Trainer, to your evaluator. See this reference for more details.