Evaluator

FunctionalEvaluator

class nni.retiarii.evaluator.FunctionalEvaluator(*args, **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.

Classification

class nni.retiarii.evaluator.pytorch.Classification(criterion=<class 'torch.nn.modules.loss.CrossEntropyLoss'>, learning_rate=0.001, weight_decay=0.0, optimizer=<class 'torch.optim.adam.Adam'>, train_dataloaders=None, val_dataloaders=None, export_onnx=True, train_dataloader=None, **trainer_kwargs)

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.

• 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 = 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.retiarii.evaluator.pytorch.ClassificationModule(*args, **kwargs)

Regression

class nni.retiarii.evaluator.pytorch.Regression(criterion=<class 'torch.nn.modules.loss.MSELoss'>, learning_rate=0.001, weight_decay=0.0, optimizer=<class 'torch.optim.adam.Adam'>, train_dataloaders=None, val_dataloaders=None, export_onnx=True, train_dataloader=None, **trainer_kwargs)

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.

• 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.retiarii.evaluator.pytorch.RegressionModule(*args, **kwargs)

Utilities

class nni.retiarii.evaluator.pytorch.Trainer(*args, **kwargs)

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

class nni.retiarii.evaluator.pytorch.DataLoader(*args, **kwargs)

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

Customization

class nni.retiarii.Evaluator

Evaluator of a model. An evaluator should define where the training code is, and the configuration of training code. The configuration includes basic runtime information trainer needs to know (such as number of GPUs) or tune-able parameters (such as learning rate), depending on the implementation of training code.

Each config should define how it is interpreted in _execute(), taking only one argument which is the mutated model class. For example, functional evaluator might directly import the function and call the function.

evaluate(model_cls)

To run evaluation of a model. The model could be either a concrete model or a callable returning a model.

The concrete implementation of evaluate depends on the implementation of _execute() in sub-class.

class nni.retiarii.evaluator.pytorch.Lightning(*args, **kwargs)

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 (Optional[Any]) – 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.

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

class nni.retiarii.evaluator.pytorch.LightningModule(*args, **kwargs)

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

running_mode: Literal['multi', 'oneshot'] = 'multi'

An indicator of whether current module is running in a multi-trial experiment or an one-shot. This flag should be automatically set by experiments when they start to run.

set_model(model)

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

Parameters:

model (callable or nn.Module) – Can be a callable returning nn.Module or nn.Module.

Cross-graph Optimization (experimental)

class nni.retiarii.evaluator.pytorch.cgo.evaluator.MultiModelSupervisedLearningModule(criterion, metrics, learning_rate=0.001, weight_decay=0.0, optimizer=<class 'torch.optim.adam.Adam'>)

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

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

class nni.retiarii.evaluator.pytorch.cgo.evaluator.Classification(criterion=<class 'torch.nn.modules.loss.CrossEntropyLoss'>, learning_rate=0.001, weight_decay=0.0, optimizer=<class 'torch.optim.adam.Adam'>, train_dataloader=None, val_dataloaders=None, **trainer_kwargs)

Trainer that is used for classification.

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

class nni.retiarii.evaluator.pytorch.cgo.evaluator.Regression(criterion=<class 'torch.nn.modules.loss.MSELoss'>, learning_rate=0.001, weight_decay=0.0, optimizer=<class 'torch.optim.adam.Adam'>, train_dataloader=None, val_dataloaders=None, **trainer_kwargs)

Trainer that is used for regression.

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