重要提示

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

日志记录和检查点#

NeMo 2.0 中的三个主要类负责配置日志记录和检查点目录

nemo.lightning.pytorch.callbacks.model_checkpoint.ModelCheckpoint

处理确定何时保存检查点的逻辑。此外，此类还提供执行异步检查点的能力。

nemo.lightning.nemo_logger.NeMoLogger

负责设置日志记录目录和（可选）配置训练器的日志记录器。

nemo.lightning.resume.AutoResume

设置检查点目录并确定是否存在要从中恢复的现有检查点。

ModelCheckpoint#

NeMo 2.0 中的 ModelCheckpoint 回调是 Pytorch Lightning 的 ModelCheckpoint 的包装器。它管理训练期间何时保存和清理检查点。此外，它还支持在训练结束时保存检查点，并为异步检查点提供必要的支持。

以下是如何实例化 ModelCheckpoint 回调的示例

from nemo.lightning.pytorch.callbacks import ModelCheckpoint

checkpoint_callback = ModelCheckpoint(
    save_best_model=True,
    save_last=True,
    monitor="val_loss",
    save_top_k=2,
    every_n_train_steps=30,
    enable_nemo_ckpt_io=False,
    dirpath='my_model_directory',
)

请参阅 NeMo Lightning 和 PyTorch Lightning 的 ModelCheckpoint 类的文档，以查找受支持参数的完整列表。此处，dirpath 指的是保存检查点的目录。请注意，dirpath 是可选的。如果未提供，则默认设置为 log_dir / checkpoints，其中 log_dir 是由 NeMoLogger 确定的路径，如后续部分详细描述的那样。

此外，请注意，异步检查点是使用 MegatronStrategy 中的 ckpt_async_save 参数设置的。然后，检查点回调访问此属性以根据请求执行异步检查点。

有两种选项可将 ModelCheckpoint 回调实例传递给训练器。

  1. 将回调添加到回调集中，然后将回调直接传递给训练器

    import nemo.lightning as nl

callbacks = [checkpoint_callback]
### add any other desired callbacks...

trainer = nl.Trainer(
    ...
    callbacks = callbacks,
    ...
)

  2. 将回调传递给 NeMoLogger，如下所述。

NeMoLogger#

NeMoLogger 类负责设置 NeMo 运行的日志记录目录。记录器支持多种参数。有关详细说明，请参阅 NeMoLogger 文档。

以下是创建新的 NeMoLogger 实例的示例

from nemo.lightning import NeMoLogger
from pytorch_lightning.loggers import TensorBoardLogger

nemo_logger = NeMoLogger(
    dir='my_logging_dir',
    name='experiment1',
    use_datetime_version=False,
    tensorboard=TensorBoardLogger(
        save_dir='local_tb_path',
    )
    update_logger_directory=True,
)

默认情况下，写入日志的 log_dir 是 dir / name / version。如果未提供显式版本且 use_datetime_version 为 False，则目录将更改为 dir / name。

TensorBoard 和 WandB 记录器也可以使用 NeMoLogger 进行配置。要设置 PTL 记录器，只需初始化记录器并将其传递给 NeMoLogger 构造函数的 tensorboard 或 wandb 参数。update_logger_directory 参数控制是否更新 PTL 记录器的目录以匹配 NeMo 日志目录。如果设置为 True，PTL 记录器也会写入相同的日志目录。

如前所述，您可以选择在此处传递 ModelCheckpoint 实例，记录器将自动在您的训练器中配置检查点回调

nemo_logger = NeMoLogger(
    ...
    ckpt=checkpoint_callback,
    ...
)

一旦您的训练器被初始化，可以使用以下命令设置 NeMoLogger

nemo_logger.setup(
    trainer,
    resume_if_exists,
)

resume_if_exists 布尔值指示是否从最新的检查点恢复（如果可用）。resume_if_exists 的值应与传递到 AutoResume 中的值匹配，如下所述。

AutoResume#

AutoResume 类管理检查点路径并检查要从中恢复的现有检查点。以下是如何使用它的示例

from nemo.lightning import AutoResume

resume = AutoResume(
    resume_if_exists=True,
    resume_ignore_no_checkpoint=True,
    dirpath="checkpoint_dir_to_resume_from"
)

在脚本中，dirpath 指的是要从中恢复的检查点目录的路径。如果未提供 dirpath，则要从中恢复的目录将默认设置为 log_dir / checkpoints，其中 log_dir 由 NemoLogger 实例确定，如上一节所述。

resume_ignore_no_checkpoint 布尔值确定在 resume_if_exists 设置为 True 且在检查点目录中未找到检查点的情况下是否继续而不出错。

确保 resume_if_exists 的值与传递到 NemoLogger 实例中的参数匹配。

AutoResume 的设置方式应与 NeMoLogger 类似。

resume.setup(trainer, model)

将模型传递到设置中是可选的。仅当从 Hugging Face 或其他非 NeMo 检查点格式导入检查点时才需要它。

整合在一起#

要整合在一起，在 NeMo 2.0 中配置记录器和检查点程序如下所示

checkpoint_callback = ModelCheckpoint(
    save_best_model=True,
    save_last=True,
    monitor="reduced_train_loss",
    save_top_k=2,
    every_n_train_steps=30,
    enable_nemo_ckpt_io=False,
    dirpath='my_model_directory',
)

logger = nemo_logger = NeMoLogger(
    dir='my_logging_dir',
    name='experiment1',
    use_datetime_version=False,
    update_logger_directory=True,
    ckpt=checkpoint_callback,
)

resume = AutoResume(
    resume_if_exists=True,
    resume_ignore_no_checkpoint=True,
)

### setup your trainer here ###

nemo_logger.setup(
    trainer,
    getattr(resume, "resume_if_exists", False),
)
resume.setup(trainer)
本页内容