Evaluator

FunctionalEvaluator

class nni.nas.evaluator.FunctionalEvaluator(function, **kwargs)[source]

Functional evaluator that directly takes a function and thus should be general. See Evaluator for instructions on how to write this function.

function

The full name of the function.

arguments

Keyword arguments for the function other than model.

Classification

class nni.nas.evaluator.pytorch.Classification(*args, **kwargs)[source]

Evaluator that is used for classification.

Available callback metrics in Classification are:

  • train_loss

  • train_acc

  • val_loss

  • val_acc

Parameters:

  • 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_dataloaders (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.

  • datamodule (LightningDataModule | None) – Used in trainer.fit(). See Lightning DataModule.

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

  • num_classes (int) – Number of classes for classification task. Required for torchmetrics >= 0.11.0. default: None

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

Examples

>>> evaluator = Classification()

To use customized criterion and optimizer:

>>> evaluator = Classification(nn.LabelSmoothingCrossEntropy, optimizer=torch.optim.SGD)

Extra keyword arguments will be passed to trainer, some of which might be necessary to enable GPU acceleration:

>>> evaluator = Classification(accelerator='gpu', devices=2, strategy='ddp')
class nni.nas.evaluator.pytorch.ClassificationModule(*args, **kwargs)[source]

Regression

class nni.nas.evaluator.pytorch.Regression(*args, **kwargs)[source]

Evaluator that is used for regression.

Available callback metrics in Regression are:

  • train_loss

  • train_mse

  • val_loss

  • val_mse

Parameters:

  • 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_dataloaders (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.

  • datamodule (LightningDataModule | None) – Used in trainer.fit(). See Lightning DataModule.

  • 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.

Examples

>>> evaluator = Regression()

Extra keyword arguments will be passed to trainer, some of which might be necessary to enable GPU acceleration:

>>> evaluator = Regression(gpus=1)
class nni.nas.evaluator.pytorch.RegressionModule(*args, **kwargs)[source]

Utilities

class nni.nas.evaluator.pytorch.Trainer(*args, **kwargs)[source]

Traced version of pytorch_lightning.Trainer. See https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html

class nni.nas.evaluator.pytorch.DataLoader(*args, **kwargs)[source]

Traced version of torch.utils.data.DataLoader. See https://pytorch.org/docs/stable/data.html

Customization

class nni.nas.evaluator.Evaluator[source]

Base class of evaluator.

To users, the evaluator is to assess the quality of a model and return a score. When an evaluator is defined, it usually accepts a few arguments, such as basic runtime information (e.g., whether to use GPU), dataset used, as well as hyper-parameters (such as learning rate). These parameters can be sometimes tunable and searched by algorithms (see MutableEvaluator).

Different evaluators could have different use scenarios and requirements on the model. For example, Classification is tailored for classification models, and assumes the model has a forward method that takes a batch of data and returns logits. Evaluators might also have different assumptions, some of which are requirements of certain algorithms. The evaluator with the most freedom is FunctionalEvaluator, but it’s also incompatible with some algorithms.

To developers, the evaluator is to implement all the logics involving forward/backward of neural networks. Sometimes the algorithm requires the training and searching at the same time (e.g., one-shot algos). In that case, although the searching part doesn’t logically belong to the evaluator, it is still the evaluator’s responsibility to implement it, and the search algorithms will make sure to properly manipulate the evaluator to achieve the goal.

Tip

Inside evaluator, you can use standard NNI trial APIs to communicate with the exploration strategy. Common usages include:

  1. Use nni.get_current_parameter() to get the current ExecutableModelSpace. Notice that ExecutableModelSpace is not a directly-runnable model (e.g., a PyTorch model), which is different from the model received in evaluate(). ExecutableModelSpace objects are useful for debugging, as well as for some evaluators which need to know extra details about how the model is sampled.

  2. Use nni.report_intermediate_result() to report intermediate results.

  3. Use nni.report_final_result() to report final results.

These APIs are only available when the evaluator is executed by NNI. We recommend using nni.get_current_parameter() is not None to check if the APIs are available before using them. Please AVOID using nni.get_next_parameter() because NAS framework has already handled the logic of retrieving the next parameter. Incorrectly using nni.get_next_parameter() may cause unexpected behavior.

evaluate(model)[source]

To run evaluation of a model. The model is usually a concrete model. The return value of evaluate() can be anything. Typically it’s used for test purposes.

Subclass should override this.

static mock_runtime(model)[source]

Context manager to mock trial APIs for standalone usage.

Under the with-context of this method, nni.get_current_parameter() will return the given model.

NOTE: This method might become a utility in trial command channel in future.

Parameters:

model (ExecutableModelSpace) – The model to be evaluated. It should be a ExecutableModelSpace object.

Examples

This method should be mostly used when testing a evaluator. A typical use case is as follows:

>>> frozen_model = model_space.freeze(sample)
>>> with evaluator.mock_runtime(frozen_model):
...     evaluator.evaluate(frozen_model.executable_model())
class nni.nas.evaluator.pytorch.Lightning(*args, **kwargs)[source]

Delegate the whole training to PyTorch Lightning.

Since the arguments passed to the initialization needs to be serialized, LightningModule, Trainer or DataLoader in this file should be used. Another option is to hide dataloader in the Lightning module, in which case, dataloaders are not required for this class to work.

Following the programming style of Lightning, metrics sent to NNI should be obtained from callback_metrics in trainer. Two hooks are added at the end of validation epoch and the end of fit, respectively. The metric name and type depend on the specific task.

Warning

The Lightning evaluator are stateful. If you try to use a previous Lightning evaluator, please note that the inner lightning_module and trainer will be reused.

Parameters:

  • lightning_module (LightningModule) – Lightning module that defines the training logic.

  • trainer (Trainer) – Lightning trainer that handles the training.

  • train_dataloders – Used in trainer.fit(). A PyTorch DataLoader with training samples. If the lightning_module has a predefined train_dataloader method this will be skipped. It can be any types of dataloader supported by Lightning.

  • val_dataloaders (Any | None) – 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. It can be any types of dataloader supported by Lightning.

  • datamodule (LightningDataModule | None) – Used in trainer.fit(). See Lightning DataModule.

  • fit_kwargs (Dict[str, Any] | None) – Keyword arguments passed to trainer.fit().

  • detect_interrupt (bool) – Lightning has a graceful shutdown mechanism. It does not terminate the whole program (but only the training) when a KeyboardInterrupt is received. Setting this to True will raise the KeyboardInterrupt to the main process, so that the whole program can be terminated.

Examples

Users should define a Lightning module that inherits LightningModule, and use Trainer and DataLoader from `nni.nas.evaluator.pytorch, and make them parameters of this evaluator:

import nni
from nni.nas.evaluator.pytorch.lightning import Lightning, LightningModule, Trainer, DataLoader
class nni.nas.evaluator.pytorch.LightningModule(*args, **kwargs)[source]

Basic wrapper of generated model. Lightning modules used in NNI should inherit this class.

It’s a subclass of pytorch_lightning.LightningModule. See https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html

See SupervisedLearningModule as an example.

property model: Module

The inner model (architecture) to train / evaluate.

It will be only available after calling set_model().

set_model(model)[source]

Set the inner model (architecture) to train / evaluate.

As there is no explicit method to “unset” a model, the model is left in the lightning module after the method is called. We don’t recommend relying on this behavior.

Cross-graph Optimization (experimental)