Important

NeMo 2.0 is an experimental feature and currently released in the dev container only: nvcr.io/nvidia/nemo:dev. Please refer to the Migration Guide for information on getting started.

Checkpoints

Pre-trained SSL checkpoints available in NeMo need to be further fine-tuned on down-stream task. There are two main ways to load pretrained checkpoints in NeMo:

• Using the restore_from() method to load a local checkpoint file (.nemo), or

• Using the from_pretrained() method to download and set up a checkpoint from NGC.

Refer to the following sections for instructions and examples for each.

Note that these instructions are for fine-tuning. To resume an unfinished training experiment, use the Experiment Manager to do so by setting the resume_if_exists flag to True.

Loading Local Checkpoints

NeMo automatically saves checkpoints of a model that is trained in a .nemo format. Alternatively, to manually save the model at any point, issue model.save_to(<checkpoint_path>.nemo).

If there is a local .nemo checkpoint that you’d like to load, use the restore_from() method:

import nemo.collections.asr as nemo_asr
ssl_model = nemo_asr.models.<MODEL_BASE_CLASS>.restore_from(restore_path="<path/to/checkpoint/file.nemo>")

Where the model base class is the ASR model class of the original checkpoint, or the general ASRModel class.

Loading NGC Pretrained Checkpoints

The SSL collection has checkpoints of several models trained on various datasets. These checkpoints are obtainable via NGC NeMo Automatic Speech Recognition collection. The model cards on NGC contain more information about each of the checkpoints available.

The table at the end of this page lists the SSL models available from NGC. The models can be accessed via the from_pretrained() method inside the ASR Model class. In general, you can load any of these models with code in the following format:

import nemo.collections.asr as nemo_asr
ssl_model = nemo_asr.models.ASRModel.from_pretrained(model_name="<MODEL_NAME>")

Where the model_name is the value under “Model Name” entry in the tables below.

For example, to load the conformer Large SSL checkpoint, run:

ssl_model = nemo_asr.models.ASRModel.from_pretrained(model_name="ssl_en_conformer_large")

You can also call from_pretrained() from the specific model class (such as SpeechEncDecSelfSupervisedModel for Conformer) if you need to access a specific model functionality.

If you would like to programatically list the models available for a particular base class, you can use the list_available_models() method.

nemo_asr.models.<MODEL_BASE_CLASS>.list_available_models()

Loading SSL checkpoint into Down-stream Model

After loading an SSL checkpoint as shown above, it’s state_dict needs to be copied to a down-stream model for fine-tuning.

For example, to load a SSL checkpoint for ASR down-stream task using EncDecRNNTBPEModel, run:

# define down-stream model
asr_model = nemo_asr.models.EncDecRNNTBPEModel(cfg=cfg.model, trainer=trainer)

# load ssl checkpoint
asr_model.load_state_dict(ssl_model.state_dict(), strict=False)

# discard ssl model
del ssl model

Refer to SSL configs to do this automatically via config files.

Fine-tuning on Downstream Datasets

After loading SSL checkpoint into down-stream model, refer to multiple ASR tutorials provided in the Tutorials section. Most of these tutorials explain how to fine-tune on some dataset as a demonstration.

Inference Execution Flow Diagram

When preparing your own inference scripts after downstream fine-tuning, please follow the execution flow diagram order for correct inference, found at the examples directory for ASR collection.

SSL Models

Below is a list of all the SSL models that are available in NeMo.

Model Name	Model Base Class	Model Card
ssl_en_conformer_large	SpeechEncDecSelfSupervisedModel	https://ngc.nvidia.com/catalog/models/nvidia:nemo:ssl_en_conformer_large
ssl_en_conformer_xlarge	SpeechEncDecSelfSupervisedModel	https://ngc.nvidia.com/catalog/models/nvidia:nemo:ssl_en_conformer_xlarge