重要

您正在查看 NeMo 2.0 文档。此版本对 API 和一个新的库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 中的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

检查点#

在 NeMo 中加载预训练检查点主要有三种方法

  • 使用 restore_from() 方法加载本地检查点文件 (.nemo)，或者

  • 将部分训练的 .ckpt（中间）检查点转换为 .nemo 格式。

  • 将 HuggingFace 公共检查点转换为 .nemo 格式。

请参阅以下部分，获取每种方法的说明和示例。

请注意，这些说明用于加载完全训练的检查点以进行评估或微调。

加载 .nemo 检查点#

NeMo 会自动保存以 .nemo 格式训练的模型的检查点。或者，要在任何时候手动保存模型，请执行 model.save_to(<checkpoint_path>.nemo)。

如果有一个本地 .nemo 检查点您想要加载，请使用 restore_from() 方法

import nemo.collections.multimodal as nemo_multimodal
model = nemo_multimodal.models.<MODEL_BASE_CLASS>.restore_from(restore_path="<path/to/checkpoint/file.nemo>")

其中模型基类是原始检查点的 MM 模型类。

转换中间检查点#

要评估部分训练的检查点，您可能需要将其转换为 .nemo 格式。转换检查点的脚本 <ADD convert_ckpt_to_nemo.py PATH>。

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>

转换 HuggingFace 检查点#

为了充分利用 NeMo 的优化训练管道和框架/TRT 推理管道，我们提供了脚本将 HuggingFace 上的常用检查点转换为 NeMo 格式。转换后，您可以对这些检查点执行微调或推理。

Stable Diffusion & ControlNet#

我们提供了 将 Huggingface 检查点转换为 .nemo 格式的脚本 <ADD convert_hf_ckpt_to_nemo.py PATH>，然后可以在我们的推理管道中使用。

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：配置文件，用于与模型权重结合生成 .nemo 检查点。它可以从虚拟运行生成，例如，可以在 nemo_experiments/stable-diffusion-train/version_0/hparams.yaml 找到。

  • model_type：在此脚本中，我们支持转换 stable_diffusion 和 controlnet 检查点。

  • nemo_clip_path：仅当 hparams_file 中的 cond_stage_config 指向 NeMo CLIP 模型时才需要。当 cond_stage_config 指向 Hugginface CLIP 时，它将被忽略。有关更多详细信息，请参阅 模型配置。

Imagen#

如果 Imagen 研究团队将来发布他们的检查点，我们将提供转换脚本。DeepFloyd IF 模型的转换脚本将在下一个版本中提供。

本页内容