Retiarii API Reference

Inline Mutation APIs

class nni.retiarii.nn.pytorch.LayerChoice(*args, **kwargs)[源代码]

Layer choice selects one of the candidates, then apply it on inputs and return results.

Layer choice does not allow itself to be nested.

参数

  • candidates (list of nn.Module or OrderedDict) -- A module list to be selected from.

  • prior (list of float) -- Prior distribution used in random sampling.

  • label (str) -- Identifier of the layer choice.

length

Deprecated. Number of ops to choose from. len(layer_choice) is recommended.

Type

int

names

Names of candidates.

Type

list of str

choices

Deprecated. A list of all candidate modules in the layer choice module. list(layer_choice) is recommended, which will serve the same purpose.

Type

list of Module

提示

candidates can be a list of modules or a ordered dict of named modules, for example,

self.op_choice = LayerChoice(OrderedDict([
    ("conv3x3", nn.Conv2d(3, 16, 128)),
    ("conv5x5", nn.Conv2d(5, 16, 128)),
    ("conv7x7", nn.Conv2d(7, 16, 128))
]))

Elements in layer choice can be modified or deleted. Use del self.op_choice["conv5x5"] or self.op_choice[1] = nn.Conv3d(...). Adding more choices is not supported yet.

classmethod create_fixed_module(candidates: Union[Dict[str, torch.nn.modules.module.Module], List[torch.nn.modules.module.Module]], *, label: Optional[str] = None, **kwargs)[源代码]

Try to create a fixed module from fixed dict. If the code is running in a trial, this method would succeed, and a concrete module instead of a mutable will be created. Raises no context error if the creation failed.

forward(x)[源代码]

Defines the computation performed at every call.

Should be overridden by all subclasses.

备注

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class nni.retiarii.nn.pytorch.InputChoice(*args, **kwargs)[源代码]

Input choice selects n_chosen inputs from choose_from (contains n_candidates keys). Use reduction to specify how chosen inputs are reduced into one output. A few options are:

  • none: do nothing and return the list directly.

  • sum: summing all the chosen inputs.

  • mean: taking the average of all chosen inputs.

  • concat: concatenate all chosen inputs at dimension 1.

We don't support customizing reduction yet.

参数

  • n_candidates (int) -- Number of inputs to choose from. It is required.

  • n_chosen (int) -- Recommended inputs to choose. If None, mutator is instructed to select any.

  • reduction (str) -- mean, concat, sum or none.

  • prior (list of float) -- Prior distribution used in random sampling.

  • label (str) -- Identifier of the input choice.

classmethod create_fixed_module(n_candidates: int, n_chosen: Optional[int] = 1, reduction: str = 'sum', *, prior: Optional[List[float]] = None, label: Optional[str] = None, **kwargs)[源代码]

Try to create a fixed module from fixed dict. If the code is running in a trial, this method would succeed, and a concrete module instead of a mutable will be created. Raises no context error if the creation failed.

forward(candidate_inputs: List[torch.Tensor]) → torch.Tensor[源代码]

Defines the computation performed at every call.

Should be overridden by all subclasses.

备注

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class nni.retiarii.nn.pytorch.ValueChoice(*args, **kwargs)[源代码]

ValueChoice is to choose one from candidates.

In most use scenarios, ValueChoice should be passed to the init parameters of a serializable module. For example,

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, nn.ValueChoice([32, 64]), kernel_size=nn.ValueChoice([3, 5, 7]))

    def forward(self, x):
        return self.conv(x)

In case, you want to search a parameter that is used repeatedly, this is also possible by sharing the same value choice instance. (Sharing the label should have the same effect.) For example,

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        hidden_dim = nn.ValueChoice([128, 512])
        self.fc = nn.Sequential(
            nn.Linear(64, hidden_dim),
            nn.Linear(hidden_dim, 10)
        )

        # the following code has the same effect.
        # self.fc = nn.Sequential(
        #     nn.Linear(64, nn.ValueChoice([128, 512], label='dim')),
        #     nn.Linear(nn.ValueChoice([128, 512], label='dim'), 10)
        # )

    def forward(self, x):
        return self.fc(x)

Note that ValueChoice should be used directly. Transformations like nn.Linear(32, nn.ValueChoice([64, 128]) * 2) are not supported.

Another common use case is to initialize the values to choose from in init and call the module in forward to get the chosen value. Usually, this is used to pass a mutable value to a functional API like torch.xxx or nn.functional.xxx`. For example,

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.dropout_rate = nn.ValueChoice([0., 1.])

    def forward(self, x):
        return F.dropout(x, self.dropout_rate())
参数

  • candidates (list) -- List of values to choose from.

  • prior (list of float) -- Prior distribution to sample from.

  • label (str) -- Identifier of the value choice.

classmethod create_fixed_module(candidates: List[Any], *, label: Optional[str] = None, **kwargs)[源代码]

Try to create a fixed module from fixed dict. If the code is running in a trial, this method would succeed, and a concrete module instead of a mutable will be created. Raises no context error if the creation failed.

forward()[源代码]

Defines the computation performed at every call.

Should be overridden by all subclasses.

备注

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class nni.retiarii.nn.pytorch.ChosenInputs(chosen: Union[List[int], int], reduction: str)[源代码]

A module that chooses from a tensor list and outputs a reduced tensor. The already-chosen version of InputChoice.

forward(candidate_inputs)[源代码]

Defines the computation performed at every call.

Should be overridden by all subclasses.

备注

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class nni.retiarii.nn.pytorch.Repeat(*args, **kwargs)[源代码]

Repeat a block by a variable number of times.

参数

  • blocks (function, list of function, module or list of module) -- The block to be repeated. If not a list, it will be replicated into a list. If a list, it should be of length max_depth, the modules will be instantiated in order and a prefix will be taken. If a function, it will be called (the argument is the index) to instantiate a module. Otherwise the module will be deep-copied.

  • depth (int or tuple of int) -- If one number, the block will be repeated by a fixed number of times. If a tuple, it should be (min, max), meaning that the block will be repeated at least min times and at most max times.

classmethod create_fixed_module(blocks: Union[Callable[[int], torch.nn.modules.module.Module], List[Callable[[int], torch.nn.modules.module.Module]], torch.nn.modules.module.Module, List[torch.nn.modules.module.Module]], depth: Union[int, Tuple[int, int]], *, label: Optional[str] = None)[源代码]

Try to create a fixed module from fixed dict. If the code is running in a trial, this method would succeed, and a concrete module instead of a mutable will be created. Raises no context error if the creation failed.

forward(x)[源代码]

Defines the computation performed at every call.

Should be overridden by all subclasses.

备注

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class nni.retiarii.nn.pytorch.Cell(op_candidates: Union[Callable, List[torch.nn.modules.module.Module]], num_nodes: int, num_ops_per_node: int = 1, num_predecessors: int = 1, merge_op: str = 'all', label: Optional[str] = None)[源代码]

Cell structure [zophnas] [zophnasnet] that is popularly used in NAS literature.

A cell consists of multiple "nodes". Each node is a sum of multiple operators. Each operator is chosen from op_candidates, and takes one input from previous nodes and predecessors. Predecessor means the input of cell. The output of cell is the concatenation of some of the nodes in the cell (currently all the nodes).

参数

  • op_candidates (function or list of module) -- A list of modules to choose from, or a function that returns a list of modules.

  • num_nodes (int) -- Number of nodes in the cell.

  • num_ops_per_node (int) -- Number of operators in each node. The output of each node is the sum of all operators in the node. Default: 1.

  • num_predecessors (int) -- Number of inputs of the cell. The input to forward should be a list of tensors. Default: 1.

  • merge_op (str) -- Currently only all is supported, which has slight difference with that described in reference. Default: all.

  • label (str) -- Identifier of the cell. Cell sharing the same label will semantically share the same choice.

引用

zophnas

Barret Zoph, Quoc V. Le, "Neural Architecture Search with Reinforcement Learning". https://arxiv.org/abs/1611.01578

zophnasnet

Barret Zoph, Vijay Vasudevan, Jonathon Shlens, Quoc V. Le, "Learning Transferable Architectures for Scalable Image Recognition". https://arxiv.org/abs/1707.07012

forward(x: List[torch.Tensor])[源代码]

Defines the computation performed at every call.

Should be overridden by all subclasses.

备注

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Graph Mutation APIs

class nni.retiarii.Mutator(sampler: Optional[nni.retiarii.mutator.Sampler] = None, label: Optional[str] = None)[源代码]

Mutates graphs in model to generate new model. Mutator class will be used in two places:

  1. Inherit Mutator to implement graph mutation logic.

  2. Use Mutator subclass to implement NAS strategy.

In scenario 1, the subclass should implement Mutator.mutate() interface with Mutator.choice(). In scenario 2, strategy should use constructor or Mutator.bind_sampler() to initialize subclass, and then use Mutator.apply() to mutate model. For certain mutator subclasses, strategy or sampler can use Mutator.dry_run() to predict choice candidates. # Method names are open for discussion.

If mutator has a label, in most cases, it means that this mutator is applied to nodes with this label.

apply(model: nni.retiarii.graph.Model) → nni.retiarii.graph.Model[源代码]

Apply this mutator on a model. Returns mutated model. The model will be copied before mutation and the original model will not be modified.

bind_sampler(sampler: nni.retiarii.mutator.Sampler) → nni.retiarii.mutator.Mutator[源代码]

Set the sampler which will handle Mutator.choice calls.

choice(candidates: Iterable[Any]) → Any[源代码]

Ask sampler to make a choice.

dry_run(model: nni.retiarii.graph.Model) → Tuple[List[List[Any]], nni.retiarii.graph.Model][源代码]

Dry run mutator on a model to collect choice candidates. If you invoke this method multiple times on same or different models, it may or may not return identical results, depending on how the subclass implements Mutator.mutate().

mutate(model: nni.retiarii.graph.Model) → None[源代码]

Abstract method to be implemented by subclass. Mutate a model in place.

class nni.retiarii.Model(_internal=False)[源代码]

Represents a neural network model.

During mutation, one Model object is created for each trainable snapshot. For example, consider a mutator that insert a node at an edge for each iteration. In one iteration, the mutator invokes 4 primitives: add node, remove edge, add edge to head, add edge to tail. These 4 primitives operates in one Model object. When they are all done the model will be set to "frozen" (trainable) status and be submitted to execution engine. And then a new iteration starts, and a new Model object is created by forking last model.

python_class

Python class that base model is converted from.

python_init_params

Initialization parameters of python class.

status

See ModelStatus.

root_graph

The outermost graph which usually takes dataset as input and feeds output to loss function.

graphs

All graphs (subgraphs) in this model.

evaluator

Model evaluator

history

Mutation history. self is directly mutated from self.history[-1]; self.history[-1] is mutated from `self.history[-2], and so on. self.history[0] is the base graph.

metric

Training result of the model, or None if it's not yet trained or has failed to train.

intermediate_metrics

Intermediate training metrics. If the model is not trained, it's an empty list.

fork() → nni.retiarii.graph.Model[源代码]

Create a new model which has same topology, names, and IDs to current one.

Can only be invoked on a frozen model. The new model will be in Mutating state.

This API is used in mutator base class.

get_node_by_name(node_name: str) → nni.retiarii.graph.Node[源代码]

Traverse all the nodes to find the matched node with the given name.

get_node_by_python_name(python_name: str) → nni.retiarii.graph.Node[源代码]

Traverse all the nodes to find the matched node with the given python_name.

get_nodes() → Iterable[nni.retiarii.graph.Node][源代码]

Traverse through all the nodes.

get_nodes_by_label(label: str) → List[nni.retiarii.graph.Node][源代码]

Traverse all the nodes to find the matched node(s) with the given label. There could be multiple nodes with the same label. Name space name can uniquely identify a graph or node.

NOTE: the implementation does not support the class abstraction

get_nodes_by_type(type_name: str) → List[nni.retiarii.graph.Node][源代码]

Traverse all the nodes to find the matched node(s) with the given type.

class nni.retiarii.Graph(model: nni.retiarii.graph.Model, graph_id: int, name: Optional[str] = None, _internal: bool = False)[源代码]

Graph topology.

This class simply represents the topology, with no semantic meaning. All other information like metric, non-graph functions, mutation history, etc should go to Model.

Each graph belongs to and only belongs to one Model.

model

The model containing (and owning) this graph.

id

Unique ID in the model. If two models have graphs of identical ID, they are semantically the same graph. Typically this means one graph is mutated from another, or they are both mutated from one ancestor.

name

Mnemonic name of this graph. It should have an one-to-one mapping with ID.

input_names

Optional mnemonic names of input parameters.

output_names

Optional mnemonic names of output values.

input_node

...

output_node

...

hidden_nodes

...

nodes

All input/output/hidden nodes.

edges

...

python_name

The name of torch.nn.Module, should have one-to-one mapping with items in python model.

fork() → nni.retiarii.graph.Graph[源代码]

Fork the model and returns corresponding graph in new model. This shortcut might be helpful because many algorithms only cares about "stem" subgraph instead of whole model.

get_node_by_id(node_id: int) → Optional[nni.retiarii.graph.Node][源代码]

Returns the node which has specified name; or returns None if no node has this name.

get_node_by_name(name: str) → Optional[nni.retiarii.graph.Node][源代码]

Returns the node which has specified name; or returns None if no node has this name.

get_node_by_python_name(python_name: str) → Optional[nni.retiarii.graph.Node][源代码]

Returns the node which has specified python_name; or returns None if no node has this python_name.

get_nodes_by_type(operation_type: str) → List[nni.retiarii.graph.Node][源代码]

Returns nodes whose operation is specified typed.

class nni.retiarii.Node(graph, node_id, name, operation, _internal=False)[源代码]

An operation or an opaque subgraph inside a graph.

Each node belongs to and only belongs to one Graph. Nodes should never be created with constructor. Use Graph.add_node() instead.

The node itself is for topology only. Information of tensor calculation should all go inside operation attribute.

TODO: parameter of subgraph (cell) It's easy to assign parameters on cell node, but it's hard to "use" them. We need to design a way to reference stored cell parameters in inner node operations. e.g. self.fc = Linear(self.units) <- how to express self.units in IR?

graph

The graph containing this node.

id

Unique ID in the model. If two models have nodes with same ID, they are semantically the same node.

name

Mnemonic name. It should have an one-to-one mapping with ID.

python_name

The name of torch.nn.Module, should have one-to-one mapping with items in python model.

label

Optional. If two nodes have the same label, they are considered same by the mutator.

operation

...

cell

Read only shortcut to get the referenced subgraph. If this node is not a subgraph (is a primitive operation), accessing cell will raise an error.

predecessors

Predecessor nodes of this node in the graph. This is an optional mutation helper.

successors

Successor nodes of this node in the graph. This is an optional mutation helper.

incoming_edges

Incoming edges of this node in the graph. This is an optional mutation helper.

outgoing_edges

Outgoing edges of this node in the graph. This is an optional mutation helper.

specialize_cell() → nni.retiarii.graph.Graph[源代码]

Only available if the operation is a cell. Duplicate the cell template and let this node reference to newly created copy.

class nni.retiarii.Edge(head: Tuple[nni.retiarii.graph.Node, Optional[int]], tail: Tuple[nni.retiarii.graph.Node, Optional[int]], _internal: bool = False)[源代码]

A tensor, or "data flow", between two nodes.

Example forward code snippet: ` a, b, c = split(x) p = concat(a, c) q = sum(b, p) z = relu(q) `

Edges in above snippet:

  • head: (split, 0), tail: (concat, 0) # a in concat

  • head: (split, 2), tail: (concat, 1) # c in concat

  • head: (split, 1), tail: (sum, -1 or 0) # b in sum

  • head: (concat, null), tail: (sum, -1 or 1) # p in sum

  • head: (sum, null), tail: (relu, null) # q in relu

graph

...

head

Head node.

tail

Tail node.

head_slot

Index of outputs in head node. If the node has only one output, this should be null.

tail_slot

Index of inputs in tail node. If the node has only one input, this should be null. If the node does not care about order, this can be -1.

class nni.retiarii.Operation(type_name: str, parameters: Dict[str, Any] = {}, _internal: bool = False, attributes: Dict[str, Any] = {})[源代码]

Calculation logic of a graph node.

The constructor is private. Use Operation.new() to create operation object.

Operation is a naive record. Do not "mutate" its attributes or store information relate to specific node. All complex logic should be implemented in Node class.

type

Operation type name (e.g. Conv2D). If it starts with underscore, the "operation" is a special one (e.g. subgraph, input/output).

parameters

Arbitrary key-value parameters (e.g. kernel_size).

Evaluators

class nni.retiarii.evaluator.FunctionalEvaluator(function, **kwargs)[源代码]

Functional evaluator that directly takes a function and thus should be general.

function

The full name of the function.

arguments

Keyword arguments for the function other than model.

class nni.retiarii.evaluator.pytorch.lightning.LightningModule(*args: Any, **kwargs: Any)[源代码]

Basic wrapper of generated model.

Lightning modules used in NNI should inherit this class.

class nni.retiarii.evaluator.pytorch.lightning.Classification(criterion: torch.nn.modules.module.Module = <class 'torch.nn.modules.loss.CrossEntropyLoss'>, learning_rate: float = 0.001, weight_decay: float = 0.0, optimizer: torch.optim.optimizer.Optimizer = <class 'torch.optim.adam.Adam'>, train_dataloader: typing.Optional[torch.utils.data.dataloader.DataLoader] = None, val_dataloaders: typing.Optional[typing.Union[torch.utils.data.dataloader.DataLoader, typing.List[torch.utils.data.dataloader.DataLoader]]] = None, export_onnx: bool = True, **trainer_kwargs)[源代码]

Trainer that is used for classification.

参数

  • criterion (nn.Module) -- Class for criterion module (not an instance). default: nn.CrossEntropyLoss

  • learning_rate (float) -- Learning rate. default: 0.001

  • weight_decay (float) -- L2 weight decay. default: 0

  • optimizer (Optimizer) -- Class for optimizer (not an instance). default: Adam

  • train_dataloders (DataLoader) -- Used in trainer.fit(). A PyTorch DataLoader with training samples. If the lightning_module has a predefined train_dataloader method this will be skipped.

  • val_dataloaders (DataLoader or List of DataLoader) -- Used in trainer.fit(). Either a single PyTorch Dataloader or a list of them, specifying validation samples. If the lightning_module has a predefined val_dataloaders method this will be skipped.

  • export_onnx (bool) -- If true, model will be exported to model.onnx before training starts. default true

  • trainer_kwargs (dict) -- Optional keyword arguments passed to trainer. See Lightning documentation for details.

class nni.retiarii.evaluator.pytorch.lightning.Regression(criterion: torch.nn.modules.module.Module = <class 'torch.nn.modules.loss.MSELoss'>, learning_rate: float = 0.001, weight_decay: float = 0.0, optimizer: torch.optim.optimizer.Optimizer = <class 'torch.optim.adam.Adam'>, train_dataloader: typing.Optional[torch.utils.data.dataloader.DataLoader] = None, val_dataloaders: typing.Optional[typing.Union[torch.utils.data.dataloader.DataLoader, typing.List[torch.utils.data.dataloader.DataLoader]]] = None, export_onnx: bool = True, **trainer_kwargs)[源代码]

Trainer that is used for regression.

参数

  • criterion (nn.Module) -- Class for criterion module (not an instance). default: nn.MSELoss

  • learning_rate (float) -- Learning rate. default: 0.001

  • weight_decay (float) -- L2 weight decay. default: 0

  • optimizer (Optimizer) -- Class for optimizer (not an instance). default: Adam

  • train_dataloders (DataLoader) -- Used in trainer.fit(). A PyTorch DataLoader with training samples. If the lightning_module has a predefined train_dataloader method this will be skipped.

  • val_dataloaders (DataLoader or List of DataLoader) -- Used in trainer.fit(). Either a single PyTorch Dataloader or a list of them, specifying validation samples. If the lightning_module has a predefined val_dataloaders method this will be skipped.

  • export_onnx (bool) -- If true, model will be exported to model.onnx before training starts. default: true

  • trainer_kwargs (dict) -- Optional keyword arguments passed to trainer. See Lightning documentation for details.

Oneshot Trainers

class nni.retiarii.oneshot.pytorch.DartsTrainer(model, loss, metrics, optimizer, num_epochs, dataset, grad_clip=5.0, learning_rate=0.0025, batch_size=64, workers=4, device=None, log_frequency=None, arc_learning_rate=0.0003, unrolled=False)[源代码]

DARTS trainer.

参数

  • model (nn.Module) -- PyTorch model to be trained.

  • loss (callable) -- Receives logits and ground truth label, return a loss tensor.

  • metrics (callable) -- Receives logits and ground truth label, return a dict of metrics.

  • optimizer (Optimizer) -- The optimizer used for optimizing the model.

  • num_epochs (int) -- Number of epochs planned for training.

  • dataset (Dataset) -- Dataset for training. Will be split for training weights and architecture weights.

  • grad_clip (float) -- Gradient clipping. Set to 0 to disable. Default: 5.

  • learning_rate (float) -- Learning rate to optimize the model.

  • batch_size (int) -- Batch size.

  • workers (int) -- Workers for data loading.

  • device (torch.device) -- torch.device("cpu") or torch.device("cuda").

  • log_frequency (int) -- Step count per logging.

  • arc_learning_rate (float) -- Learning rate of architecture parameters.

  • unrolled (float) -- True if using second order optimization, else first order optimization.

class nni.retiarii.oneshot.pytorch.EnasTrainer(model, loss, metrics, reward_function, optimizer, num_epochs, dataset, batch_size=64, workers=4, device=None, log_frequency=None, grad_clip=5.0, entropy_weight=0.0001, skip_weight=0.8, baseline_decay=0.999, ctrl_lr=0.00035, ctrl_steps_aggregate=20, ctrl_kwargs=None)[源代码]

ENAS trainer.

参数

  • model (nn.Module) -- PyTorch model to be trained.

  • loss (callable) -- Receives logits and ground truth label, return a loss tensor.

  • metrics (callable) -- Receives logits and ground truth label, return a dict of metrics.

  • reward_function (callable) -- Receives logits and ground truth label, return a tensor, which will be feeded to RL controller as reward.

  • optimizer (Optimizer) -- The optimizer used for optimizing the model.

  • num_epochs (int) -- Number of epochs planned for training.

  • dataset (Dataset) -- Dataset for training. Will be split for training weights and architecture weights.

  • batch_size (int) -- Batch size.

  • workers (int) -- Workers for data loading.

  • device (torch.device) -- torch.device("cpu") or torch.device("cuda").

  • log_frequency (int) -- Step count per logging.

  • grad_clip (float) -- Gradient clipping. Set to 0 to disable. Default: 5.

  • entropy_weight (float) -- Weight of sample entropy loss.

  • skip_weight (float) -- Weight of skip penalty loss.

  • baseline_decay (float) -- Decay factor of baseline. New baseline will be equal to baseline_decay * baseline_old + reward * (1 - baseline_decay).

  • ctrl_lr (float) -- Learning rate for RL controller.

  • ctrl_steps_aggregate (int) -- Number of steps that will be aggregated into one mini-batch for RL controller.

  • ctrl_steps (int) -- Number of mini-batches for each epoch of RL controller learning.

  • ctrl_kwargs (dict) -- Optional kwargs that will be passed to ReinforceController.

class nni.retiarii.oneshot.pytorch.ProxylessTrainer(model, loss, metrics, optimizer, num_epochs, dataset, warmup_epochs=0, batch_size=64, workers=4, device=None, log_frequency=None, arc_learning_rate=0.001, grad_reg_loss_type=None, grad_reg_loss_params=None, applied_hardware=None, dummy_input=(1, 3, 224, 224), ref_latency=65.0)[源代码]

Proxyless trainer.

参数

  • model (nn.Module) -- PyTorch model to be trained.

  • loss (callable) -- Receives logits and ground truth label, return a loss tensor.

  • metrics (callable) -- Receives logits and ground truth label, return a dict of metrics.

  • optimizer (Optimizer) -- The optimizer used for optimizing the model.

  • num_epochs (int) -- Number of epochs planned for training.

  • dataset (Dataset) -- Dataset for training. Will be split for training weights and architecture weights.

  • warmup_epochs (int) -- Number of epochs to warmup model parameters.

  • batch_size (int) -- Batch size.

  • workers (int) -- Workers for data loading.

  • device (torch.device) -- torch.device("cpu") or torch.device("cuda").

  • log_frequency (int) -- Step count per logging.

  • arc_learning_rate (float) -- Learning rate of architecture parameters.

  • grad_reg_loss_type (string) -- Regularization type to add hardware related loss, allowed types include - "mul#log": regularized_loss = (torch.log(expected_latency) / math.log(self.ref_latency)) ** beta - "add#linear": regularized_loss = reg_lambda * (expected_latency - self.ref_latency) / self.ref_latency - None: do not apply loss regularization.

  • grad_reg_loss_params (dict) -- Regularization params, allowed params include - "alpha" and "beta" is required when grad_reg_loss_type == "mul#log" - "lambda" is required when grad_reg_loss_type == "add#linear"

  • applied_hardware (string) -- Applied hardware for to constraint the model's latency. Latency is predicted by Microsoft nn-Meter (https://github.com/microsoft/nn-Meter).

  • dummy_input (tuple) -- The dummy input shape when applied to the target hardware.

  • ref_latency (float) -- Reference latency value in the applied hardware (ms).

class nni.retiarii.oneshot.pytorch.SinglePathTrainer(model, loss, metrics, optimizer, num_epochs, dataset_train, dataset_valid, batch_size=64, workers=4, device=None, log_frequency=None)[源代码]

Single-path trainer. Samples a path every time and backpropagates on that path.

参数

  • model (nn.Module) -- Model with mutables.

  • loss (callable) -- Called with logits and targets. Returns a loss tensor.

  • metrics (callable) -- Returns a dict that maps metrics keys to metrics data.

  • optimizer (Optimizer) -- Optimizer that optimizes the model.

  • num_epochs (int) -- Number of epochs of training.

  • dataset_train (Dataset) -- Dataset of training.

  • dataset_valid (Dataset) -- Dataset of validation.

  • batch_size (int) -- Batch size.

  • workers (int) -- Number of threads for data preprocessing. Not used for this trainer. Maybe removed in future.

  • device (torch.device) -- Device object. Either torch.device("cuda") or torch.device("cpu"). When None, trainer will automatic detects GPU and selects GPU first.

  • log_frequency (int) -- Number of mini-batches to log metrics.

Exploration Strategies

class nni.retiarii.strategy.Random(variational=False, dedup=True, model_filter=None)[源代码]

Random search on the search space.

参数

  • variational (bool) -- Do not dry run to get the full search space. Used when the search space has variational size or candidates. Default: false.

  • dedup (bool) -- Do not try the same configuration twice. When variational is true, deduplication is not supported. Default: true.

  • model_filter (Callable[[Model], bool]) -- Feed the model and return a bool. This will filter the models in search space and select which to submit.

class nni.retiarii.strategy.GridSearch(shuffle=True)[源代码]

Traverse the search space and try all the possible combinations one by one.

参数

shuffle (bool) -- Shuffle the order in a candidate list, so that they are tried in a random order. Default: true.

class nni.retiarii.strategy.RegularizedEvolution(optimize_mode='maximize', population_size=100, sample_size=25, cycles=20000, mutation_prob=0.05, on_failure='ignore', model_filter=None)[源代码]

Algorithm for regularized evolution (i.e. aging evolution). Follows "Algorithm 1" in Real et al. "Regularized Evolution for Image Classifier Architecture Search".

参数

  • optimize_mode (str) -- Can be one of "maximize" and "minimize". Default: maximize.

  • population_size (int) -- The number of individuals to keep in the population. Default: 100.

  • cycles (int) -- The number of cycles (trials) the algorithm should run for. Default: 20000.

  • sample_size (int) -- The number of individuals that should participate in each tournament. Default: 25.

  • mutation_prob (float) -- Probability that mutation happens in each dim. Default: 0.05

  • on_failure (str) -- Can be one of "ignore" and "worst". If "ignore", simply give up the model and find a new one. If "worst", mark the model as -inf (if maximize, inf if minimize), so that the algorithm "learns" to avoid such model. Default: ignore.

  • model_filter (Callable[[Model], bool]) -- Feed the model and return a bool. This will filter the models in search space and select which to submit.

class nni.retiarii.strategy.TPEStrategy[源代码]

The Tree-structured Parzen Estimator (TPE) [bergstrahpo] is a sequential model-based optimization (SMBO) approach. SMBO methods sequentially construct models to approximate the performance of hyperparameters based on historical measurements, and then subsequently choose new hyperparameters to test based on this model.

引用

bergstrahpo

Bergstra et al., "Algorithms for Hyper-Parameter Optimization". https://papers.nips.cc/paper/4443-algorithms-for-hyper-parameter-optimization.pdf

class nni.retiarii.strategy.PolicyBasedRL(max_collect: int = 100, trial_per_collect=20, policy_fn: Optional[Callable[[nni.retiarii.strategy._rl_impl.ModelEvaluationEnv], tianshou.policy.base.BasePolicy]] = None)[源代码]

Algorithm for policy-based reinforcement learning. This is a wrapper of algorithms provided in tianshou (PPO by default), and can be easily customized with other algorithms that inherit BasePolicy (e.g., REINFORCE 1).

参数

  • max_collect (int) -- How many times collector runs to collect trials for RL. Default 100.

  • trial_per_collect (int) -- How many trials (trajectories) each time collector collects. After each collect, trainer will sample batch from replay buffer and do the update. Default: 20.

  • policy_fn (function) -- Takes ModelEvaluationEnv as input and return a policy. See _default_policy_fn for an example.

引用

1

Barret Zoph and Quoc V. Le, "Neural Architecture Search with Reinforcement Learning". https://arxiv.org/abs/1611.01578

Retiarii Experiments

class nni.retiarii.experiment.pytorch.RetiariiExperiment(base_model: torch.nn.modules.module.Module, trainer: Union[nni.retiarii.graph.Evaluator, nni.retiarii.oneshot.interface.BaseOneShotTrainer], applied_mutators: Optional[List[nni.retiarii.mutator.Mutator]] = None, strategy: Optional[nni.retiarii.strategy.base.BaseStrategy] = None)[源代码]
export_top_models(top_k: int = 1, optimize_mode: str = 'maximize', formatter: str = 'dict') → Any[源代码]

Export several top performing models.

For one-shot algorithms, only top-1 is supported. For others, optimize_mode and formatter are available for customization.

top_kint

How many models are intended to be exported.

optimize_modestr

maximize or minimize. Not supported by one-shot algorithms. optimize_mode is likely to be removed and defined in strategy in future.

formatterstr

Support code and dict. Not supported by one-shot algorithms. If code, the python code of model will be returned. If dict, the mutation history will be returned.

retrain_model(model)[源代码]

this function retrains the exported model, and test it to output test accuracy

run(config: Optional[nni.retiarii.experiment.pytorch.RetiariiExeConfig] = None, port: int = 8080, debug: bool = False) → str[源代码]

Run the experiment. This function will block until experiment finish or error.

start(port: int = 8080, debug: bool = False) → None[源代码]

Start the experiment in background. This method will raise exception on failure. If it returns, the experiment should have been successfully started. :param port: The port of web UI. :param debug: Whether to start in debug mode.

stop() → None[源代码]

Stop background experiment.

class nni.retiarii.experiment.pytorch.RetiariiExeConfig(training_service_platform: Optional[str] = None, **kwargs)[源代码]
validate(initialized_tuner: bool = False) → None[源代码]

Validate legality of the config object. Raise exception if any error occurred.

This function does not return truth value. Do not write if config.validate().

返回类型

None

CGO Execution

nni.retiarii.evaluator.pytorch.cgo.evaluator.MultiModelSupervisedLearningModule(criterion: torch.nn.modules.module.Module, metrics: typing.Dict[str, torchmetrics.metric.Metric], n_models: int = 0, learning_rate: float = 0.001, weight_decay: float = 0.0, optimizer: torch.optim.optimizer.Optimizer = <class 'torch.optim.adam.Adam'>)[源代码]

Lightning Module of SupervisedLearning for Cross-Graph Optimization. Users who needs cross-graph optimization should use this module.

参数

  • criterion (nn.Module) -- Class for criterion module (not an instance). default: nn.CrossEntropyLoss

  • learning_rate (float) -- Learning rate. default: 0.001

  • weight_decay (float) -- L2 weight decay. default: 0

  • optimizer (Optimizer) -- Class for optimizer (not an instance). default: Adam

nni.retiarii.evaluator.pytorch.cgo.evaluator.Classification(lightning_module: nni.retiarii.evaluator.pytorch.lightning.LightningModule, trainer: pytorch_lightning.trainer.trainer.Trainer, train_dataloader: Optional[torch.utils.data.dataloader.DataLoader] = None, val_dataloaders: Optional[Union[torch.utils.data.dataloader.DataLoader, List[torch.utils.data.dataloader.DataLoader]]] = None)[源代码]

Trainer that is used for classification.

参数

  • criterion (nn.Module) -- Class for criterion module (not an instance). default: nn.CrossEntropyLoss

  • learning_rate (float) -- Learning rate. default: 0.001

  • weight_decay (float) -- L2 weight decay. default: 0

  • optimizer (Optimizer) -- Class for optimizer (not an instance). default: Adam

  • train_dataloders (DataLoader) -- Used in trainer.fit(). A PyTorch DataLoader with training samples. If the lightning_module has a predefined train_dataloader method this will be skipped.

  • val_dataloaders (DataLoader or List of DataLoader) -- Used in trainer.fit(). Either a single PyTorch Dataloader or a list of them, specifying validation samples. If the lightning_module has a predefined val_dataloaders method this will be skipped.

  • trainer_kwargs (dict) -- Optional keyword arguments passed to trainer. See Lightning documentation for details.

nni.retiarii.evaluator.pytorch.cgo.evaluator.Regression(lightning_module: nni.retiarii.evaluator.pytorch.lightning.LightningModule, trainer: pytorch_lightning.trainer.trainer.Trainer, train_dataloader: Optional[torch.utils.data.dataloader.DataLoader] = None, val_dataloaders: Optional[Union[torch.utils.data.dataloader.DataLoader, List[torch.utils.data.dataloader.DataLoader]]] = None)[源代码]

Trainer that is used for regression.

参数

  • criterion (nn.Module) -- Class for criterion module (not an instance). default: nn.MSELoss

  • learning_rate (float) -- Learning rate. default: 0.001

  • weight_decay (float) -- L2 weight decay. default: 0

  • optimizer (Optimizer) -- Class for optimizer (not an instance). default: Adam

  • train_dataloders (DataLoader) -- Used in trainer.fit(). A PyTorch DataLoader with training samples. If the lightning_module has a predefined train_dataloader method this will be skipped.

  • val_dataloaders (DataLoader or List of DataLoader) -- Used in trainer.fit(). Either a single PyTorch Dataloader or a list of them, specifying validation samples. If the lightning_module has a predefined val_dataloaders method this will be skipped.

  • trainer_kwargs (dict) -- Optional keyword arguments passed to trainer. See Lightning documentation for details.

Utilities

nni.retiarii.basic_unit(cls: nni.retiarii.serializer.T, basic_unit_tag: bool = True) → Union[nni.retiarii.serializer.T, nni.common.serializer.Traceable][源代码]

To wrap a module as a basic unit, is to make it a primitive and stop the engine from digging deeper into it.

basic_unit_tag is true by default. If set to false, it will not be explicitly mark as a basic unit, and graph parser will continue to parse. Currently, this is to handle a special case in nn.Sequential.

Although basic_unit calls trace in its implementation, it is not for serialization. Rather, it is meant to capture the initialization arguments for mutation. Also, graph execution engine will stop digging into the inner modules when it reaches a module that is decorated with basic_unit.

@basic_unit
class PrimitiveOp(nn.Module):
    ...
nni.retiarii.model_wrapper(cls: nni.retiarii.serializer.T) → Union[nni.retiarii.serializer.T, nni.common.serializer.Traceable][源代码]

Wrap the base model (search space). For example,

@model_wrapper
class MyModel(nn.Module):
    ...

The wrapper serves two purposes:

  1. Capture the init parameters of python class so that it can be re-instantiated in another process.

  2. Reset uid in namespace so that the auto label counting in each model stably starts from zero.

Currently, NNI might not complain in simple cases where @model_wrapper is actually not needed. But in future, we might enforce @model_wrapper to be required for base model.

nni.retiarii.fixed_arch(fixed_arch: Union[str, pathlib.Path, Dict[str, Any]], verbose=True)[源代码]

Load architecture from fixed_arch and apply to model. This should be used as a context manager. For example,

with fixed_arch('/path/to/export.json'):
    model = Model(3, 224, 224)
参数

  • fixed_arc (str, Path or dict) -- Path to the JSON that stores the architecture, or dict that stores the exported architecture.

  • verbose (bool) -- Print log messages if set to True

返回

Context manager that provides a fixed architecture when creates the model.

返回类型

ContextStack