Checkpoints

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 loading fully trained checkpoints for evaluation or fine-tuning. For resuming 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
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.

NGC Pretrained Checkpoints

The ASR collection has checkpoints of several models trained on various datasets for a variety of tasks. 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 tables below list the ASR 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
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 base English QuartzNet model for speech recognition, run:

model = nemo_asr.models.ASRModel.from_pretrained(model_name="QuartzNet15x5Base-En")

You can also call from_pretrained() from the specific model class (such as EncDecCTCModel for QuartzNet) 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()

Transcribing/Inference

To perform inference and transcribe a sample of speech after loading the model, use the transcribe() method:

model.transcribe(paths2audio_files=[list of audio files], batch_size=BATCH_SIZE, logprobs=False)

Setting the argument logprobs to True returns the log probabilities instead of transcriptions. For more information, see ./api.html#modules. The audio files should be 16KHz monochannel wav files.

Fine-tuning on Different Datasets

There are multiple ASR tutorials provided in the Tutorials section. Most of these tutorials explain how to instantiate a pre-trained model, prepare the model for fine-tuning on some dataset (in the same language) as a demonstration.

Automatic Speech Recognition Models

Below is a list of all the ASR models that are available in NeMo for specific languages, as well as auxiliary language models for certain languages.

Language Models for ASR

Model Name

Model Base Class

Model Card

asrlm_en_transformer_large_ls

TransformerLMModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:asrlm_en_transformer_large_ls


Speech Recognition (Languages)

English

Model Name

Model Base Class

Model Card

QuartzNet15x5Base-En

EncDecCTCModel

https://ngc.nvidia.com/catalog/models/nvidia:nemospeechmodels

stt_en_jasper10x5dr

EncDecCTCModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_jasper10x5dr

stt_en_citrinet_256

EncDecCTCBPEModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_256

stt_en_citrinet_512

EncDecCTCBPEModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_512

stt_en_citrinet_1024

EncDecCTCBPEModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_1024

stt_en_citrinet_256_gamma_0_25

EncDecCTCBPEModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_256_gamma_0_25

stt_en_citrinet_512_gamma_0_25

EncDecCTCBPEModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_512_gamma_0_25

stt_en_citrinet_1024_gamma_0_25

EncDecCTCBPEModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_1024_gamma_0_25

stt_en_conformer_ctc_small

EncDecCTCBPEModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_small

stt_en_conformer_ctc_medium

EncDecCTCBPEModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_medium

stt_en_conformer_ctc_large

EncDecCTCBPEModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large

stt_en_conformer_ctc_small_ls

EncDecCTCBPEModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_small_ls

stt_en_conformer_ctc_medium_ls

EncDecCTCBPEModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_medium_ls

stt_en_conformer_ctc_large_ls

EncDecCTCBPEModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large_ls


Mandarin

Model

Model Base Class

Model Card

stt_zh_citrinet_512

EncDecCTCBPEModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_citrinet_512


German

Model

Model Base Class

Model Card

stt_de_quartznet15x5

EncDecCTCModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_de_quartznet15x5


Polish

Model

Model Base Class

Model Card

stt_pl_quartznet15x5

EncDecCTCModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_pl_quartznet15x5


Italian

Model

Model Base Class

Model Card

stt_it_quartznet15x5

EncDecCTCModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_it_quartznet15x5


Russian

Model

Model Base Class

Model Card

stt_ru_quartznet15x5

EncDecCTCModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ru_quartznet15x5


Spanish

Model

Model Base Class

Model Card

stt_es_quartznet15x5

EncDecCTCModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_es_quartznet15x5

stt_es_citrinet_512

EncDecCTCModelBPE

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_es_citrinet_512


Catalan

Model

Model Base Class

Model Card

stt_ca_quartznet15x5

EncDecCTCModel

https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ca_quartznet15x5