Construct Space with Mutator (legacy)


This is a legacy document for NNI v2.x. This is now currently no longer maintained.

Besides the mutation primitives demonstrated in the basic tutorial, NNI provides a more general approach to express a model space, i.e., Mutator, to cover more complex model spaces. The high-level APIs are also implemented with mutator in the underlying system, which can be seen as a special case of model mutation.


Mutator and inline mutation APIs can NOT be used together.

A mutator is a piece of logic to express how to mutate a given model. Users are free to write their own mutators. Then a model space is expressed with a base model and a list of mutators. A model in the model space is sampled by applying the mutators on the base model one after another. An example is shown below.

applied_mutators = []

BlockMutator is defined by users to express how to mutate the base model.

Write a mutator

User-defined mutator should inherit nni.retiarii.Mutator class, and implement mutation logic in the member function nni.retiarii.Mutator.mutate().

from nni.retiarii import Mutator
class BlockMutator(Mutator):
  def __init__(self, target: str, candidates: List):
      super(BlockMutator, self).__init__() = target
      self.candidate_op_list = candidates

  def mutate(self, model):
    nodes = model.get_nodes_by_label(
    for node in nodes:
      chosen_op = self.choice(self.candidate_op_list)
      node.update_operation(chosen_op.type, chosen_op.params)

The input of nni.retiarii.Mutator.mutate() is graph IR (Intermediate Representation) of the base model, users can mutate the graph using the graph’s member functions (e.g., nni.retiarii.Model.get_nodes_by_label()). The mutation operations can be combined with the API self.choice, in order to express a set of possible mutations. In the above example, the node’s operation can be changed to any operation from candidate_op_list.

Use placeholder to make mutation easier: nni.retiarii.nn.pytorch.Placeholder. If you want to mutate a subgraph or node of your model, you can define a placeholder in this model to represent the subgraph or node. Then, use mutator to mutate this placeholder to make it real modules.

ph = nn.Placeholder(
  kernel_size_options=[1, 3, 5],
  n_layer_options=[1, 2, 3, 4],

label is used by mutator to identify this placeholder. The other parameters are the information that is required by mutator. They can be accessed from node.operation.parameters as a dict, it could include any information that users want to put to pass it to user defined mutator. The complete example code can be found in Mnasnet base model.

Starting an experiment is almost the same as using inline mutation APIs. The only difference is that the applied mutators should be passed to nni.retiarii.experiment.pytorch.RetiariiExperiment. Below is a simple example.

exp = RetiariiExperiment(base_model, trainer, applied_mutators, simple_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.experiment_name = 'mnasnet_search'
exp_config.trial_concurrency = 2
exp_config.max_trial_number = 10
exp_config.training_service.use_active_gpu = False, 8081)