NeMo Core API

Classes and Interfaces

class nemo.core.ModelPT(cfg: <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e11e048>, trainer: <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e11e278> = None)

Bases: sphinx.ext.autodoc.importer._MockObject, nemo.core.classes.common.Model

Interface for Pytorch-lightning based NeMo models

cfg

configure_optimizers()

classmethod extract_state_dict_from(restore_path: str, save_dir: str, split_by_module: bool = False)

Extract the state dict(s) from a provided .nemo tarfile and save it to a directory. :param restore_path: path to .nemo file from which state dict(s) should be extracted :param save_dir: directory in which the saved state dict(s) should be stored :param split_by_module: bool flag, which determins whether the output checkpoint should

be for the entire Model, or the individual module’s that comprise the Model

Example

To convert the .nemo tarfile into a single Model level PyTorch checkpoint ` state_dict = nemo.collections.asr.models.EncDecCTCModel.extract_state_dict_from('asr.nemo', './asr_ckpts) `

To restore a model from a Model level checkpoint ` model = nemo.collections.asr.models.EncDecCTCModel(cfg)  # or any other method of restoration model.load_state_dict(torch.load("./asr_ckpts/model_weights.ckpt")) `

To convert the .nemo tarfile into multiple Module level PyTorch checkpoints ``` state_dict = nemo.collections.asr.models.EncDecCTCModel.extract_state_dict_from(‘asr.nemo’, ‘./asr_ckpts,

split_by_module=True)

```

To restore a module from a Module level checkpoint ``` model = model = nemo.collections.asr.models.EncDecCTCModel(cfg) # or any other method of restoration

# load the individual components model.preprocessor.load_state_dict(torch.load(“./asr_ckpts/preprocessor.ckpt”)) model.encoder.load_state_dict(torch.load(“./asr_ckpts/encoder.ckpt”)) model.decoder.load_state_dict(torch.load(“./asr_ckpts/decoder.ckpt”)) ```

Returns:The state dict that was loaded from the original .nemo checkpoint

get_test_dataloader_prefix(dataloader_idx: int = 0) → str

Get the name of one or more data loaders, which will be prepended to all logs.

Parameters:dataloader_idx – Index of the data loader. Returns:str name of the data loader at index provided.

get_validation_dataloader_prefix(dataloader_idx: int = 0) → str

Get the name of one or more data loaders, which will be prepended to all logs.

Parameters:dataloader_idx – Index of the data loader. Returns:str name of the data loader at index provided.

classmethod load_from_checkpoint(checkpoint_path: str, *args, map_location: Union[Dict[str, str], str, <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e1860b8>, int, Callable, None] = None, hparams_file: Optional[str] = None, strict: bool = True, **kwargs)

Loads ModelPT from checkpoint, with some maintenance of restoration. For documentation, please refer to LightningModule.load_from_checkpoin() documentation.

multi_test_epoch_end(outputs: List[Dict[str, <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e194320>]], dataloader_idx: int = 0) → Optional[Dict[str, Dict[str, <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e1944a8>]]]

multi_validation_epoch_end(outputs: List[Dict[str, <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e186f98>]], dataloader_idx: int = 0) → Optional[Dict[str, Dict[str, <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e194160>]]]

num_weights

prepare_test(trainer: <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e11e278>) → bool

Helper method to check whether the model can safely be tested on a dataset after training (or loading a checkpoint).

# Usage: trainer = Trainer() if model.prepare_test(trainer):

trainer.test(model)

Returns:bool which declares the model safe to test. Provides warnings if it has to return False to guide the user.

register_artifact(config_path: str, src: str)

Register model artifacts with this function. These artifacts (files) will be included inside .nemo file when model.save_to(“mymodel.nemo”) is called.

WARNING: If you specified /example_folder/example.txt but ./example.txt exists, then ./example.txt will be used.

Parameters:

• config_path – config path where artifact is used
• src – path to the artifact

Returns:

path to be used when accessing artifact. If src=’’ or None then ‘’ or None will be returned

classmethod restore_from(restore_path: str, override_config_path: Optional[str] = None, map_location: Optional[<sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e11e470>] = None)

Restores model instance (weights and configuration) into .nemo file :param restore_path: path to .nemo file from which model should be instantiated :param override_config_path: path to a yaml config that will override the internal

config file

Parameters:

• map_location – Optional torch.device() to map the instantiated model to a device. By default (None), it will select a GPU if available, falling back to CPU otherwise.
• Example –

  ` model = nemo.collections.asr.models.EncDecCTCModel.restore_from('asr.nemo') assert isinstance(model, nemo.collections.asr.models.EncDecCTCModel) `

Returns:

An instance of type cls

save_to(save_path: str)

Saves model instance (weights and configuration) into .nemo file. You can use “restore_from” method to fully restore instance from .nemo file.

.nemo file is an archive (tar.gz) with the following:

model_config.yaml - model configuration in .yaml format. You can deserialize this into cfg argument for model’s constructor model_wights.chpt - model checkpoint

Parameters:save_path – Path to .nemo file where model instance should be saved

set_trainer(trainer: <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e11e278>)

Set an instance of Trainer object.

Parameters:trainer – PyTorch Lightning Trainer object.

setup_multiple_test_data(test_data_config: Union[<sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e11e048>, Dict[KT, VT]])

(Optionally) Setups data loader to be used in test, with support for multiple data loaders.

Parameters:test_data_layer_config – test data layer parameters.

setup_multiple_validation_data(val_data_config: Union[<sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e11e048>, Dict[KT, VT]])

(Optionally) Setups data loader to be used in validation, with support for multiple data loaders.

Parameters:val_data_layer_config – validation data layer parameters.

setup_optimization(optim_config: Union[<sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e11e048>, Dict[KT, VT], None] = None)

Prepares an optimizer from a string name and its optional config parameters.

Parameters:optim_config –

A dictionary containing the following keys:

• ”lr”: mandatory key for learning rate. Will raise ValueError if not provided.
• ”optimizer”: string name pointing to one of the available optimizers in the registry. If not provided, defaults to “adam”.
• ”opt_args”: Optional list of strings, in the format “arg_name=arg_value”. The list of “arg_value” will be parsed and a dictionary of optimizer kwargs will be built and supplied to instantiate the optimizer.

setup_test_data(test_data_config: Union[<sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e11e048>, Dict[KT, VT]])

(Optionally) Setups data loader to be used in test

Parameters:test_data_layer_config – test data layer parameters.

Returns:

setup_training_data(train_data_config: Union[<sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e11e048>, Dict[KT, VT]])

Setups data loader to be used in training

Parameters:train_data_layer_config – training data layer parameters.

Returns:

setup_validation_data(val_data_config: Union[<sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e11e048>, Dict[KT, VT]])

(Optionally) Setups data loader to be used in validation :param val_data_layer_config: validation data layer parameters.

Returns:

teardown(stage: str)

Called at the end of fit and test.

Parameters:stage – either ‘fit’ or ‘test’

test_dataloader()

test_epoch_end(outputs: Union[List[Dict[str, <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e186978>]], List[List[Dict[str, <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e186b00>]]]]) → Optional[Dict[str, Dict[str, <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e186dd8>]]]

Default DataLoader for Test set which automatically supports multiple data loaders via multi_test_epoch_end.

If multi dataset support is not required, override this method entirely in base class. In such a case, there is no need to implement multi_test_epoch_end either.

Note

If more than one data loader exists, and they all provide test_loss, only the test_loss of the first data loader will be used by default. This default can be changed by passing the special key test_loss_idx: int inside the test_ds config.

Parameters:outputs – Single or nested list of tensor outputs from one or more data loaders. Returns:A dictionary containing the union of all items from individual data_loaders, along with merged logs from all data loaders.

test_step(batch, batch_ix)

train_dataloader()

training_step(batch, batch_ix)

val_dataloader()

validation_epoch_end(outputs: Union[List[Dict[str, <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e186358>]], List[List[Dict[str, <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e1864e0>]]]]) → Optional[Dict[str, Dict[str, <sphinx.ext.autodoc.importer._MockObject object at 0x7fa31e1867b8>]]]

Default DataLoader for Validation set which automatically supports multiple data loaders via multi_validation_epoch_end.

If multi dataset support is not required, override this method entirely in base class. In such a case, there is no need to implement multi_validation_epoch_end either.

Note

If more than one data loader exists, and they all provide val_loss, only the val_loss of the first data loader will be used by default. This default can be changed by passing the special key val_loss_idx: int inside the validation_ds config.

Parameters:outputs – Single or nested list of tensor outputs from one or more data loaders. Returns:A dictionary containing the union of all items from individual data_loaders, along with merged logs from all data loaders.

validation_step(batch, batch_ix)

Neural Types

class nemo.core.neural_types.neural_type.NeuralType(axes: Optional[Tuple] = None, elements_type: nemo.core.neural_types.elements.ElementType = VoidType, optional=False)

Bases: object

This is the main class which would represent neural type concept. It is used to represent the types of inputs and outputs. :param axes: a tuple of AxisTypes objects representing the semantics of what varying each axis means

You can use a short, string-based form here. For example: (‘B’, ‘C’, ‘H’, ‘W’) would correspond to an NCHW format frequently used in computer vision. (‘B’, ‘T’, ‘D’) is frequently used for signal processing and means [batch, time, dimension/channel].

Parameters:

• elements_type (ElementType) – an instance of ElementType class representing the semantics of what is stored inside the tensor. For example: logits (LogitsType), log probabilities (LogprobType), etc.
• optional (bool) – By default, this is false. If set to True, it would means that input to the port of this type can be optional.

compare(second) → nemo.core.neural_types.comparison.NeuralTypeComparisonResult

Performs neural type comparison of self with second. When you chain two modules’ inputs/outputs via __call__ method, this comparison will be called to ensure neural type compatibility.

compare_and_raise_error(parent_type_name, port_name, second_object)

Method compares definition of one type with another and raises an error if not compatible.

exception nemo.core.neural_types.neural_type.NeuralTypeError

Bases: Exception

Base class for neural type related exceptions.

exception nemo.core.neural_types.neural_type.NeuralPortNameMismatchError(input_port_name)

Bases: nemo.core.neural_types.neural_type.NeuralTypeError

Exception raised when neural module is called with incorrect port names.

exception nemo.core.neural_types.neural_type.NeuralPortNmTensorMismatchError(class_name, port_name, first_type, second_type, type_comatibility)

Bases: nemo.core.neural_types.neural_type.NeuralTypeError

Exception raised when a port is fed with a NmTensor of incompatible type.