| Checkpoints | |
| =========== | |
| There are three main ways to load pretrained checkpoints in NeMo: | |
| * Using the :code:`restore_from()` method to load a local checkpoint file (``.nemo``), or | |
| * Converting a partially trained ``.ckpt`` (intermediate) checkpoint to ``.nemo`` format. | |
| * Converting HuggingFace public checkpoints to ``.nemo`` format. | |
| 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. | |
| Loading ``.nemo`` 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 :code:`model.save_to(<checkpoint_path>.nemo)`. | |
| If there is a local ``.nemo`` checkpoint that you'd like to load, use the :code:`restore_from()` method: | |
| .. code-block:: python | |
| import nemo.collections.multimodal as nemo_multimodal | |
| model = nemo_multimodal.models.<MODEL_BASE_CLASS>.restore_from(restore_path="<path/to/checkpoint/file.nemo>") | |
| Where the model base class is the MM model class of the original checkpoint. | |
| Converting Intermediate Checkpoints | |
| ----------------------------------- | |
| To evaluate a partially trained checkpoint, you may need to convert it to ``.nemo`` format. | |
| `script to convert the checkpoint <ADD convert_ckpt_to_nemo.py PATH>`. | |
| .. code-block:: python | |
| python -m torch.distributed.launch --nproc_per_node=<tensor_model_parallel_size> * <pipeline_model_parallel_size> \ | |
| convert_ckpt_to_nemo.py \ | |
| --checkpoint_folder <path_to_PTL_checkpoints_folder> \ | |
| --checkpoint_name <checkpoint_name> \ | |
| --nemo_file_path <path_to_output_nemo_file> \ | |
| --tensor_model_parallel_size <tensor_model_parallel_size> \ | |
| --pipeline_model_parallel_size <pipeline_model_parallel_size> | |
| Converting HuggingFace Checkpoints | |
| ---------------------------------- | |
| To fully utilize the optimized training pipeline and framework/TRT inference pipeline | |
| of NeMo, we provide scripts to convert popular checkpoints on HuggingFace into NeMo format. | |
| Once converted, you can perform fine-tuning or inference on such checkpoints. | |
| Stable Diffusion & ControlNet | |
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ | |
| We provide `script to convert the Huggingface checkpoint <ADD convert_hf_ckpt_to_nemo.py PATH>` to ``.nemo`` format, which can then be used within our inference pipeline. | |
| .. code-block:: python | |
| python convert_hf_ckpt_to_nemo.py \ | |
| --ckpt_path <path_to_HF_checkpoints> \ | |
| --hparams_file <path_to_hparams_file> \ | |
| --nemo_file_path <path_to_output_nemo_file> \ | |
| --model_type <model_to_be_converted> \ | |
| --nemo_clip_path <clip_ckpt_in_nemo_format> | |
| - ``hparams_file``: Config file to be combined with model weights to generate ``.nemo`` checkpoint. It can be generated from a dummy run and can be found at, for example, ``nemo_experiments/stable-diffusion-train/version_0/hparams.yaml``. | |
| - ``model_type``: We support converting `stable_diffusion` and `controlnet` checkpoint in this script. | |
| - ``nemo_clip_path``: It's required only when the ``cond_stage_config`` in ``hparams_file`` refer to a NeMo CLIP model. It will be ignored when ``cond_stage_config`` refer to Hugginface CLIP. See :ref:`sd-config-section` for more details. | |
| Imagen | |
| ^^^^^^^^^^^^^^ | |
| We will provide conversion script if Imagen research team releases their checkpoint | |
| in the future. Conversion script for DeepFloyd IF models will be provided in the | |
| next release. | |