重要提示

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

NeMo Core API#

所有 NeMo 模型的基础类#

class nemo.core.ModelPT(*args: Any, **kwargs: Any)#

基类： LightningModule, Model

基于 Pytorch-lightning 的 NeMo 模型的接口

on_fit_start() → None#
register_artifact(
config_path: str,
src: str,
verify_src_exists: bool = True,
)#

使用此函数注册模型工件。当调用 model.save_to(“mymodel.nemo”) 时，这些工件（文件）将包含在 .nemo 文件中。

工作原理

  1. 它始终返回现有的绝对路径，该路径可在 Model 构造函数调用期间使用

    异常：src 为 None 或 “”，在这种情况下，将不执行任何操作，并且将返回 src

  2. 它将 (config_path, model_utils.ArtifactItem()) 对添加到 self.artifacts

    If "src" is local existing path:
    then it will be returned in absolute path form.
elif "src" starts with "nemo_file:unique_artifact_name":
    .nemo will be untarred to a temporary folder location and an actual existing path will be returned
else:
    an error will be raised.

警告：在模型的构造函数中使用 .register_artifact 调用。返回的路径不能保证在您退出模型的构造函数后仍然存在。

参数:

  • config_path (str) – 工件键。通常对应于模型配置。

  • src (str) – 工件的路径。

  • verify_src_exists (bool) – 如果设置为 False，则工件是可选的，即使找不到 src，register_artifact 也会返回 None。默认为 True。

返回值:

如果 src 不是 None 或空字符串，它始终返回绝对路径，保证在模型实例生命周期内存在

返回类型:

str

has_artifacts() → bool#

如果模型已注册工件，则返回 True

has_native_or_submodules_artifacts() → bool#

如果它有工件或任何子模块有工件，则返回 True

has_nemo_submodules() → bool#

如果它有任何已注册的 NeMo 子模块，则返回 True

register_nemo_submodule(
name: str,
config_field: str,
model: ModelPT,
) → None#

将 NeMo 模型添加为子模块。子模块可以通过注册此子模块的父 NeMo 模型（self）上的 name 属性访问。在保存过程中，整个父模型 (self) 被视为一个实体模型，其中包含来自子模块的工件，子模块配置将被保存到父模型的 config_field 中。此方法对于创建嵌套模型是必要的，例如：

class ParentModel(ModelPT):
    def __init__(self, cfg, trainer=None):
        super().__init__(cfg=cfg, trainer=trainer)

        # annotate type for autocompletion and type checking (optional)
        self.child_model: Optional[ChildModel] = None
        if cfg.get("child_model") is not None:
            self.register_nemo_submodule(
                name="child_model",
                config_field="child_model",
                model=ChildModel(self.cfg.child_model, trainer=trainer),
            )
        # ... other code
参数:

  • name – 子模块的属性名称

  • config_field – 配置中的字段，子模块配置应保存在此处

  • model – NeMo 模型，ModelPT 的实例

named_nemo_modules(
prefix_name: str = '',
prefix_config: str = '',
) → Iterator[Tuple[str, str, ModelPT]]#

返回递归遍历所有 NeMo 子模块的迭代器，从核心模块开始，产生 (属性路径、配置路径、子模块) 元组

参数:

  • prefix_name – 名称路径的前缀

  • prefix_config – 配置路径的前缀

返回值:

遍历 (属性路径、配置路径、子模块) 的迭代器，从 (prefix, self) 开始

save_to(save_path: str)#
将模型实例（权重和配置）保存到 .nemo 文件中

您可以使用 “restore_from” 方法从 .nemo 文件完全恢复实例。

.nemo 文件是一个存档文件 (tar.gz)，包含以下内容：

model_config.yaml - .yaml 格式的模型配置。您可以将其反序列化为模型构造函数的 cfg 参数 model_wights.ckpt - 模型检查点

参数:

save_path – 应在其中保存模型实例的 .nemo 文件的路径

classmethod restore_from(
restore_path: str,
override_config_path: omegaconf.OmegaConf | str | None = None,
map_location: torch.device | None = None,
strict: bool = True,
return_config: bool = False,
save_restore_connector: SaveRestoreConnector | None = None,
trainer: lightning.pytorch.Trainer | None = None,
validate_access_integrity: bool = True,
)#

从 .nemo 文件恢复模型实例（权重和配置）。

参数:

  • restore_path – 从中实例化模型的 .nemo 文件的路径

  • override_config_path – 将覆盖内部配置文件或表示模型配置的 OmegaConf / DictConfig 对象的 yaml 配置文件的路径。

  • map_location – 可选的 torch.device()，用于将实例化的模型映射到设备。默认情况下 (None)，如果可用，它将选择 GPU，否则回退到 CPU。

  • strict – 传递给 load_state_dict。默认为 True。

  • return_config – 如果设置为 true，则仅返回恢复模型的底层配置作为 OmegaConf DictConfig 对象，而不实例化模型。

  • trainer – 可选，pytorch lightning Trainer 对象，将转发到实例化的模型的构造函数。

  • save_restore_connector (SaveRestoreConnector) – 可以被覆盖以添加自定义保存和恢复逻辑。

  • 示例 –

    ` model = nemo.collections.asr.models.EncDecCTCModel.restore_from('asr.nemo') assert isinstance(model, nemo.collections.asr.models.EncDecCTCModel) `

返回值:

cls 类型或其底层配置的实例（如果设置了 return_config）。

classmethod load_from_checkpoint(
checkpoint_path: str,
*args,
map_location: Dict[str, str] | str | torch.device | int | Callable | None = None,
hparams_file: str | None = None,
strict: bool = True,
**kwargs,
)#

从检查点加载 ModelPT，并进行一些恢复维护。有关文档，请参阅 LightningModule.load_from_checkpoint() 文档。

abstract setup_training_data(
train_data_config: omegaconf.DictConfig | Dict,
)#

设置用于训练的数据加载器

参数:

train_data_layer_config – 训练数据层参数。

返回值

abstract setup_validation_data(
val_data_config: omegaconf.DictConfig | Dict,
)#

设置用于验证的数据加载器 :param val_data_layer_config: 验证数据层参数。

返回值

setup_test_data(
test_data_config: omegaconf.DictConfig | Dict,
)#

（可选）设置用于测试的数据加载器

参数:

test_data_layer_config – 测试数据层参数。

返回值

setup_multiple_validation_data(
val_data_config: omegaconf.DictConfig | Dict,
)#

（可选）设置用于验证的数据加载器，支持多个数据加载器。

参数:

val_data_layer_config – 验证数据层参数。

setup_multiple_test_data(
test_data_config: omegaconf.DictConfig | Dict,
)#

（可选）设置用于测试的数据加载器，支持多个数据加载器。

参数:

test_data_layer_config – 测试数据层参数。

setup_megatron_optimization(
optim_config: Dict[str, Any] | omegaconf.DictConfig,
)#

设置 mcore 优化器配置。

参数:

optim_config – 用于设置 Mcore 优化器选项的 NeMo optim 参数。

setup_optimization(
optim_config: omegaconf.DictConfig | Dict | None = None,
optim_kwargs: Dict[str, Any] | None = None,
)#

从字符串名称及其可选配置参数准备优化器。

参数:

  • optim_config –

    包含以下键的字典

    • “lr”：学习率的必填键。如果未提供，将引发 ValueError。

    • “optimizer”：指向注册表中可用优化器之一的字符串名称。如果未提供，则默认为 “adam”。

    • “opt_args”：可选的字符串列表，格式为 “arg_name=arg_value”。将解析 “arg_value” 列表，并构建优化器 kwargs 字典并提供给实例化优化器。

  • optim_kwargs – 包含优化器的其他 kwargs 的字典。用于与 OmegaConf 不兼容的非原始类型。

setup_optimizer_param_groups()#

用于为优化器创建参数组。例如，这可以用于指定每层学习率

optim.SGD([

{‘params’: model.base.parameters()}, {‘params’: model.classifier.parameters(), ‘lr’: 1e-3} ], lr=1e-2, momentum=0.9)

有关更多信息，请参阅 http://pytorch.ac.cn/docs/stable/optim.html。默认情况下，ModelPT 将使用 self.parameters()。覆盖此方法以添加自定义参数组。在配置文件中，添加 ‘optim_param_groups’ 以支持不同组件的不同 LR（未指定的参数将使用默认 LR）

模型
optim_param_groups
encoder

lr: 1e-4 momentum: 0.8

decoder

lr: 1e-3

optim

lr: 3e-3 momentum: 0.9

configure_optimizers()#
propagate_model_guid()#

将模型 GUID 递归传播到所有子模块。

setup(stage: str | None = None)#

在 fit、validate、test 或 predict 的开始时调用。当使用 DDP 时，在每个进程上调用此函数。

参数:

stage – fit、validate、test 或 predict

train_dataloader()#
val_dataloader()#
test_dataloader()#
on_validation_epoch_end(
sync_metrics: bool = False,
) → Dict[str, Dict[str, torch.Tensor]] | None#

验证集的默认数据加载器，它通过 multi_validation_epoch_end 自动支持多个数据加载器。

如果不需要多数据集支持，请完全在基类中覆盖此方法。在这种情况下，也不需要实现 multi_validation_epoch_end。

注意

如果存在多个数据加载器，并且它们都提供 val_loss，则默认情况下仅使用第一个数据加载器的 val_loss。可以通过在 validation_ds 配置中传递特殊键 val_dl_idx: int 来更改此默认设置。

参数:

outputs – 来自一个或多个数据加载器的单个或嵌套张量输出列表。

返回值:

一个字典，包含来自各个数据加载器的所有项目的并集，以及来自所有数据加载器的合并日志。

on_test_epoch_end() → Dict[str, Dict[str, torch.Tensor]] | None#

测试集的默认数据加载器，它通过 multi_test_epoch_end 自动支持多个数据加载器。

如果不需要多数据集支持，请完全在基类中覆盖此方法。在这种情况下，也不需要实现 multi_test_epoch_end。

注意

如果存在多个数据加载器，并且它们都提供 test_loss，则默认情况下仅使用第一个数据加载器的 test_loss。可以通过在 test_ds 配置中传递特殊键 test_dl_idx: int 来更改此默认设置。

参数:

outputs – 来自一个或多个数据加载器的单个或嵌套张量输出列表。

返回值:

一个字典，包含来自各个数据加载器的所有项目的并集，以及来自所有数据加载器的合并日志。

multi_validation_epoch_end(
outputs: List[Dict[str, torch.Tensor]],
dataloader_idx: int = 0,
) → Dict[str, Dict[str, torch.Tensor]] | None#

添加对多个验证数据集的支持。应由子类覆盖，以便获得每个数据加载器的适当日志。

参数:

  • outputs – 与 LightningModule.on_validation_epoch_end() 为单个数据加载器提供的输出相同。

  • dataloader_idx – 表示数据加载器索引的整数。

返回值:

一个值字典，可以选择包含子字典 log，以便日志中的值将以数据加载器前缀开头。

multi_test_epoch_end(
outputs: List[Dict[str, torch.Tensor]],
dataloader_idx: int = 0,
) → Dict[str, Dict[str, torch.Tensor]] | None#

添加对多个测试数据集的支持。应由子类覆盖，以便获得每个数据加载器的适当日志。

参数:

  • outputs – 与 LightningModule.on_validation_epoch_end() 为单个数据加载器提供的输出相同。

  • dataloader_idx – 表示数据加载器索引的整数。

返回值:

一个值字典，可以选择包含子字典 log，以便日志中的值将以数据加载器前缀开头。

get_validation_dataloader_prefix(
dataloader_idx: int = 0,
) → str#

获取一个或多个数据加载器的名称，这些名称将添加到所有日志的前面。

参数:

dataloader_idx – 数据加载器的索引。

返回值:

str，表示所提供索引处的数据加载器的名称。

get_test_dataloader_prefix(dataloader_idx: int = 0) → str#

获取一个或多个数据加载器的名称，这些名称将添加到所有日志的前面。

参数:

dataloader_idx – 数据加载器的索引。

返回值:

str，表示所提供索引处的数据加载器的名称。

load_part_of_state_dict(
state_dict,
include,
exclude,
load_from_string=None,
)#
maybe_init_from_pretrained_checkpoint(
cfg: omegaconf.OmegaConf,
map_location: str = 'cpu',
)#

使用通过特定配置参数获得的参数初始化给定模型。所提供模型的 state dict 将使用 strict=False 设置进行更新，以防止需要精确的模型参数匹配。

初始化方式

init_from_nemo_model: .nemo 模型的字符串路径，用于从单个 nemo 文件加载 state_dict；如果从多个文件加载，则传入一个字典，其中值具有以下字段

path: .nemo 模型的字符串路径

include: 可选的字符串列表，其中至少有一个字符串需要包含在参数名称中，才能从此 .nemo 文件加载。默认值：包含所有内容。

exclude: 可选的字符串列表，可用于排除任何包含这些字符串之一的参数，使其不从此 .nemo 文件加载。默认值：不排除任何内容。

Hydra 用法示例

init_from_nemo_model
model0

path:<path/to/model1> include:[“encoder”]

model1

path:<path/to/model2> include:[“decoder”] exclude:[“embed”]

init_from_pretrained_model: 预训练模型检查点的字符串名称（通过云端获取）。

模型将被下载（或使用缓存副本），实例化，然后提取其 state dict。如果从多个模型加载，您可以传入与 init_from_nemo_model 格式相同的字典，但使用 “name” 而不是 “path”

init_from_ptl_ckpt: Pytorch Lightning 检查点文件的字符串名称。它将被加载，并且

state dict 将被提取。如果从多个文件加载，您可以传入与 init_from_nemo_model 格式相同的字典。

参数:

  • cfg – 用于实例化模型的配置。它只需要包含上述键之一。

  • map_location – str 或 torch.device()，表示预训练模型或检查点中的中间 state dict 将加载到哪里。

classmethod extract_state_dict_from(
restore_path: str,
save_dir: str,
split_by_module: bool = False,
save_restore_connector: SaveRestoreConnector | None = None,
)#

从提供的 .nemo tar 文件中提取 state dict，并将其保存到目录中。

参数:

  • restore_path – 要从中提取 state dict 的 .nemo 文件路径

  • save_dir – 应在其中存储已保存 state dict 的目录

  • split_by_module – 布尔标志，用于确定输出检查点应是整个模型，还是组成模型的各个模块的检查点

  • save_restore_connector (SaveRestoreConnector) – 可以被覆盖以添加自定义的保存和恢复逻辑。

示例

要将 .nemo tar 文件转换为单个模型级别的 PyTorch 检查点 :: state_dict = nemo.collections.asr.models.EncDecCTCModel.extract_state_dict_from(‘asr.nemo’, ‘./asr_ckpts’)

要从模型级别检查点恢复模型 :: model = nemo.collections.asr.models.EncDecCTCModel(cfg) # 或任何其他恢复方法 model.load_state_dict(torch.load(“./asr_ckpts/model_weights.ckpt”))

要将 .nemo tar 文件转换为多个模块级别的 PyTorch 检查点 :: state_dict = nemo.collections.asr.models.EncDecCTCModel.extract_state_dict_from(‘asr.nemo’, ‘./asr_ckpts’, split_by_module=True)

要从模块级别检查点恢复模块 :: model = nemo.collections.asr.models.EncDecCTCModel(cfg) # 或任何其他恢复方法

# 加载各个组件 model.preprocessor.load_state_dict(torch.load(“./asr_ckpts/preprocessor.ckpt”)) model.encoder.load_state_dict(torch.load(“./asr_ckpts/encoder.ckpt”)) model.decoder.load_state_dict(torch.load(“./asr_ckpts/decoder.ckpt”))

返回值:

从原始 .nemo 检查点加载的 state dict

prepare_test(trainer: lightning.pytorch.Trainer) → bool#

辅助方法，用于检查模型在训练后（或加载检查点后）是否可以安全地在数据集上进行测试。

trainer = Trainer()
if model.prepare_test(trainer):
    trainer.test(model)
返回值:

布尔值，声明模型可以安全地进行测试。如果必须返回 False，则提供警告以指导用户。

set_trainer(trainer: lightning.pytorch.Trainer)#

设置 Trainer 对象的实例。

参数:

trainer – PyTorch Lightning Trainer 对象。

set_world_size(trainer: lightning.pytorch.Trainer)#

从 PyTorch Lightning Trainer 确定世界大小。然后更新 AppState。

参数:

trainer (Trainer) – PyTorch Lightning Trainer 对象

summarize(
max_depth: int = 1,
) → lightning.pytorch.utilities.model_summary.ModelSummary#

总结此 LightningModule。

参数:

max_depth – 摘要将包含的最大层级嵌套深度。值为 0 将关闭层级摘要。默认值：1。

返回值:

模型摘要对象

property num_weights#

实用工具属性，返回模型的参数总数。

trainer()#
property cfg#

保存模型最终内部配置的属性。

注意

对此配置的更改不会反映在模型的状态中。请使用更新后的配置创建一个新模型，以正确更新模型。

property hparams#

覆盖默认的 hparams 属性以返回最新的模型配置。如果没有此更改，如果直接更改 self._cfg（例如，在 self.setup_optimization() 中）而不是通过 self.cfg = new_cfg 完成更改，则 hparams 属性将返回旧配置。

property validation_step_outputs#

validation_step 的缓存输出。它可以是项目列表（对于单个数据加载器）或列表的列表（对于多个数据加载器）。

返回值:

validation_step 的输出列表。

property test_step_outputs#

test_step 的缓存输出。它可以是项目列表（对于单个数据加载器）或列表的列表（对于多个数据加载器）。

返回值:

test_step 的输出列表。

classmethod update_save_restore_connector(save_restore_connector)#
on_train_start()#

PyTorch Lightning hook: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-train-start 我们在此处使用它来复制动态冻结的相关配置。

on_train_batch_start(
batch: Any,
batch_idx: int,
unused: int = 0,
) → int | None#

PyTorch Lightning hook: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-train-batch-start 我们在此处使用它来启用性能分析和动态冻结。

on_train_batch_end(
outputs,
batch: Any,
batch_idx: int,
unused: int = 0,
) → None#

PyTorch Lightning hook: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-train-batch-end 我们在此处使用它来启用 nsys 性能分析。

on_train_end()#

PyTorch Lightning hook: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-train-end 我们在此处使用它来清理动态冻结配置。

on_test_end()#

PyTorch Lightning hook: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-test-end

on_predict_end()#

PyTorch Lightning hook: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-test-end

cuda(device=None)#
PTL 正在覆盖此方法并更改模块的 pytorch 行为。

如果 device 为 None，PTL LightingModule 覆盖将模块移动到设备 0。请参阅此处的 PTL 方法： Lightning-AI/lightning

我们在此处覆盖此方法是为了保持默认的 Pytorch nn.module 行为： pytorch/pytorch

将所有模型参数和缓冲区移动到 GPU。

这也使得关联的参数和缓冲区成为不同的对象。因此，如果模块将在 GPU 上进行优化，则应在构造优化器之前调用此方法。

注意

此方法就地修改模块。

参数:

device (int, optional) – 如果指定，所有参数都将复制到该设备

返回值:

self

返回类型:

Module

基础神经模块类#

class nemo.core.NeuralModule(*args: Any, **kwargs: Any)#

基类： Module, Typing, Serialization, FileIO

抽象类，提供所有 PyTorch 神经模块之间共享的接口。

property num_weights#

实用工具属性，返回 NeuralModule 的参数总数。

input_example(max_batch=None, max_dim=None)#

如果随机输入不起作用，请覆盖此方法 :returns: 有效输入数据的元组示例。

freeze() → None#

冻结所有参数以进行推理。

此方法将模块的所有参数的 requires_grad 设置为 False。它还将每个参数的原始 requires_grad 状态存储在字典中，以便在 unfreeze() 中设置 partial=True 时，unfreeze() 可以恢复原始状态。

unfreeze(partial: bool = False) → None#

解冻所有参数以进行训练。

允许完全解冻或部分解冻（如果模块之前已使用 freeze() 显式冻结）。partial 参数用于确定是解冻所有参数，还是仅解冻之前在 freeze() 之前解冻的参数。

示例

考虑一个具有编码器和解码器模块的模型。假设我们希望编码器始终保持冻结。

`python model.encoder.freeze()  # 显式冻结编码器中的所有参数 `

在推理期间，应冻结模型的所有参数 - 我们通过调用模型的 freeze 方法来实现这一点。此步骤记录编码器模块参数已冻结，因此如果调用部分解冻，我们应保持编码器参数冻结。

`python model.freeze()  # 冻结模型中的所有参数；编码器保持冻结 `

现在，在微调期间，我们希望解冻解码器但保持编码器冻结。我们可以通过调用 unfreeze(partial=True) 来实现这一点。

`python model.unfreeze(partial=True)  # 仅解冻解码器；编码器保持冻结 `

参数:

partial – 如果为 True，则仅解冻先前冻结的参数。如果在调用 freeze() 时参数已冻结，则在调用 unfreeze(partial=True) 后，它将保持冻结状态。

as_frozen()#

上下文管理器，它临时冻结模块，产生控制权，最后部分解冻模块以返回原始状态。

允许完全解冻或部分解冻（如果模块之前已使用 freeze() 显式冻结）。partial 参数用于确定是解冻所有参数，还是仅解冻之前在 freeze() 之前解冻的参数。

示例

with model.as_frozen(): # 默认情况下，partial = True

# 对模型执行某些操作 pass

# 模型的参数现在恢复到 requires_grad 的原始状态

基础 Mixin 类#

class nemo.core.Typing

基类： ABC

赋予模块神经类型的接口

property input_types: Dict[str, NeuralType] | None

定义这些以启用输入神经类型检查

property output_types: Dict[str, NeuralType] | None

定义这些以启用输出神经类型检查

_validate_input_types(
input_types=None,
ignore_collections=False,
**kwargs,
)

此函数执行以下几项操作。

  1. 它确保 len(self.input_types <非可选>) <= len(kwargs) <= len(self.input_types)。

  2. 对于作为输入传递给包装函数的每个（关键字名称，关键字值）

    • 检查关键字名称是否存在于有效的 self.input_types 名称列表中。

    • 检查关键字值是否具有 neural_type 属性。
      • 如果它有，则执行比较检查并断言神经类型

        是兼容的（相同或更大）。

    • 检查关键字值是否为容器类型（列表或元组）。如果是，

      然后在嵌套结构的每个元素上递归执行上述神经类型的元素级测试。

参数:

  • input_types – 类级别定义的 input_types，或本地函数覆盖的类型定义。

  • ignore_collections – 为了向后兼容，可以使用此标志显式禁用容器支持。设置为 True 时，所有嵌套都将被忽略，并且跳过嵌套深度检查。

  • kwargs – argument_name:argument_value 对的字典，在调用时传递给包装函数。

_attach_and_validate_output_types(
out_objects,
ignore_collections=False,
output_types=None,
)

此函数执行以下几项操作。

  1. 它确保 len(out_object) == len(self.output_types)。

  2. 如果输出是张量（或张量的列表/元组...的列表/元组），它

    将神经类型附加到它。对于没有 neural_type 属性的对象，例如 python 对象（字典和列表、原始数据类型、结构），不附加神经类型。

注意：tensor.neural_type 仅在 _validate_input_types 期间检查，该检查在 forward() 之前调用。

参数:

  • output_types – 类级别定义的 output_types，或本地函数覆盖的类型定义。

  • ignore_collections – 为了向后兼容，可以使用此标志显式禁用容器支持。设置为 True 时，所有嵌套都将被忽略，并且跳过嵌套深度检查。

  • out_objects – 包装函数的输出。

__check_neural_type(
obj,
metadata: TypecheckMetadata,
depth: int,
name: str | None = None,
)

递归测试 obj 是否满足语义神经类型断言。如果提供形状信息，则可以包括形状检查。

参数:

  • obj – 可以分配值的任何 python 对象。

  • metadata – TypecheckMetadata 对象。

  • depth – 当前递归深度。

  • name – 源 obj 的可选名称，在发生错误时使用。

__attach_neural_type(
obj,
metadata: TypecheckMetadata,
depth: int,
name: str | None = None,
)

递归地将神经类型附加到给定的对象 - 只要它可以被分配一些值。

参数:

  • obj – 可以分配值的任何 python 对象。

  • metadata – TypecheckMetadata 对象。

  • depth – 当前递归深度。

  • name – 源 obj 的可选名称，在发生错误时使用。

class nemo.core.Serialization

基类： ABC

classmethod from_config_dict(
config: DictConfig,
trainer: Trainer | None = None,
)

使用基于 DictConfig 的配置实例化对象

to_config_dict() → omegaconf.DictConfig

将对象的配置返回到配置字典

class nemo.core.FileIO

基类： ABC

save_to(save_path: str)

用于保存包含检查点、配置和任何其他工件的 tar 文件的标准化方法。通过 nemo.core.connectors.save_restore_connector.SaveRestoreConnector.save_to() 实现。

参数:

save_path – str，文件应保存到的路径。

classmethod restore_from(
restore_path: str,
override_config_path: str | None = None,
map_location: torch.device | None = None,
strict: bool = True,
return_config: bool = False,
trainer: Trainer | None = None,
save_restore_connector: SaveRestoreConnector = None,
)

从 .nemo 文件恢复模型实例（权重和配置）

参数:

  • restore_path – 从中实例化模型的 .nemo 文件的路径

  • override_config_path – 将覆盖内部配置文件或表示模型配置的 OmegaConf / DictConfig 对象的 yaml 配置文件的路径。

  • map_location – 可选的 torch.device()，用于将实例化的模型映射到设备。默认情况下 (None)，如果可用，它将选择 GPU，否则回退到 CPU。

  • strict – 传递给 load_state_dict。默认为 True

  • return_config – 如果设置为 true，则仅返回恢复模型的底层配置作为 OmegaConf DictConfig 对象，而不实例化模型。

  • trainer – 可选的 Trainer 对象，传递给模型构造函数。

  • save_restore_connector – 可选的 SaveRestoreConnector 对象，用于定义 restore_from() 方法的实现。

classmethod from_config_file(path2yaml_file: str)

从 YAML 配置文件实例化 NeMo Model 的实例。权重将被随机初始化。 :param path2yaml_file: 包含模型配置的 yaml 文件的路径

返回值

to_config_file(path2yaml_file: str)

将当前实例的配置保存到 YAML 配置文件。权重将不会被保存。 :param path2yaml_file: path2yaml_file: 模型配置将被保存到的 yaml 文件的路径

返回值

基础连接器类#

class nemo.core.connectors.save_restore_connector.SaveRestoreConnector#

基类： object

save_to(
model: ModelPT,
save_path: str,
)#

将模型实例（权重和配置）保存到 .nemo 文件中。您可以使用 “restore_from” 方法从 .nemo 文件完全恢复实例。

.nemo 文件是一个存档文件 (tar.gz)，包含以下内容：

model_config.yaml - .yaml 格式的模型配置。您可以将其反序列化为模型构造函数的 cfg 参数 model_wights.ckpt - 模型检查点

参数:

  • model – 要保存的 ModelPT 对象。

  • save_path – 应在其中保存模型实例的 .nemo 文件的路径

返回值:

模型实例保存到的 .nemo 文件路径（与 save_path 参数相同），如果不是 rank 0，则为 None

如果标志 pack_nemo_file 设置为 False，则路径可以是目录。

返回类型:

str

load_config_and_state_dict(
calling_cls,
restore_path: str,
override_config_path: omegaconf.OmegaConf | str | None = None,
map_location: torch.device | None = None,
strict: bool = True,
return_config: bool = False,
trainer: lightning.pytorch.trainer.trainer.Trainer | None = None,
validate_access_integrity: bool = True,
)#

将模型实例（权重和配置）恢复到 .nemo 文件中

参数:

  • restore_path – 从中实例化模型的 .nemo 文件的路径

  • override_config_path – 将覆盖内部配置文件或表示模型配置的 OmegaConf / DictConfig 对象的 yaml 配置文件的路径。

  • map_location – 可选的 torch.device()，用于将实例化的模型映射到设备。默认情况下 (None)，如果可用，它将选择 GPU，否则回退到 CPU。

  • strict – 传递给 load_state_dict。默认为 True

  • return_config – 如果设置为 true，则仅返回恢复模型的底层配置作为 OmegaConf DictConfig 对象，而不实例化模型。

示例

` model = nemo.collections.asr.models.EncDecCTCModel.restore_from('asr.nemo') assert isinstance(model, nemo.collections.asr.models.EncDecCTCModel) `

返回值:

cls 类型或其底层配置的实例（如果设置了 return_config）。

modify_state_dict(conf, state_dict)#

实用方法，允许在将参数加载到模型之前修改状态字典。 :param conf: 模型级别的 OmegaConf 对象。:param state_dict: 从检查点恢复的状态字典。

返回值:

一个可能被修改的状态字典。

load_instance_with_state_dict(
instance,
state_dict,
strict,
)#

实用方法，用于加载带有（可能已修改的）状态字典的模型实例。

参数:

  • instance – ModelPT 子类实例。

  • state_dict – 状态字典（可能已被修改）

  • strict – Bool，指示在加载状态字典时是否执行严格检查。

restore_from(
calling_cls,
restore_path: str,
override_config_path: omegaconf.OmegaConf | str | None = None,
map_location: torch.device | None = None,
strict: bool = True,
return_config: bool = False,
trainer: lightning.pytorch.trainer.trainer.Trainer | None = None,
validate_access_integrity: bool = True,
)#

将模型实例（权重和配置）恢复到 .nemo 文件中

参数:

  • restore_path – 从中实例化模型的 .nemo 文件的路径

  • override_config_path – 将覆盖内部配置文件或表示模型配置的 OmegaConf / DictConfig 对象的 yaml 配置文件的路径。

  • map_location – 可选的 torch.device()，用于将实例化的模型映射到设备。默认情况下 (None)，如果可用，它将选择 GPU，否则回退到 CPU。

  • strict – 传递给 load_state_dict。默认为 True

  • return_config – 如果设置为 true，则仅返回恢复模型的底层配置作为 OmegaConf DictConfig 对象，而不实例化模型。

  • trainer – 可选的 Trainer 对象，传递给模型构造函数。

示例

` model = nemo.collections.asr.models.EncDecCTCModel.restore_from('asr.nemo') assert isinstance(model, nemo.collections.asr.models.EncDecCTCModel) `

返回值:

cls 类型或其底层配置的实例（如果设置了 return_config）。

extract_state_dict_from(
restore_path: str,
save_dir: str,
split_by_module: bool = False,
)#

从提供的 .nemo tar 文件中提取 state dict，并将其保存到目录中。

参数:

  • restore_path – 要从中提取 state dict 的 .nemo 文件路径

  • save_dir – 应在其中存储已保存 state dict 的目录

  • split_by_module – 布尔标志，用于确定输出检查点应是整个模型，还是组成模型的各个模块的检查点

示例

要将 .nemo tar 文件转换为单个模型级别的 PyTorch 检查点 :: state_dict = nemo.collections.asr.models.EncDecCTCModel.extract_state_dict_from(‘asr.nemo’, ‘./asr_ckpts’)

要从模型级别检查点恢复模型 :: model = nemo.collections.asr.models.EncDecCTCModel(cfg) # 或任何其他恢复方法 model.load_state_dict(torch.load(“./asr_ckpts/model_weights.ckpt”))

要将 .nemo tar 文件转换为多个模块级别的 PyTorch 检查点 :: state_dict = nemo.collections.asr.models.EncDecCTCModel.extract_state_dict_from(‘asr.nemo’, ‘./asr_ckpts’, split_by_module=True)

要从模块级别检查点恢复模块 :: model = nemo.collections.asr.models.EncDecCTCModel(cfg) # 或任何其他恢复方法

# 加载各个组件 model.preprocessor.load_state_dict(torch.load(“./asr_ckpts/preprocessor.ckpt”)) model.encoder.load_state_dict(torch.load(“./asr_ckpts/encoder.ckpt”)) model.decoder.load_state_dict(torch.load(“./asr_ckpts/decoder.ckpt”))

返回值:

从原始 .nemo 检查点加载的 state dict

register_artifact(
模型,
config_path: str,
src: str,
verify_src_exists: bool = True,
)#

使用此函数注册模型工件。当调用 model.save_to(“mymodel.nemo”) 时，这些工件（文件）将包含在 .nemo 文件中。

工作原理

  1. 它始终返回现有的绝对路径，该路径可在 Model 构造函数调用期间使用

    异常：src 为 None 或 “”，在这种情况下，将不执行任何操作，并且将返回 src

  2. 它将 (config_path, model_utils.ArtifactItem()) 对添加到 self.artifacts

    If "src" is local existing path:
    then it will be returned in absolute path form
elif "src" starts with "nemo_file:unique_artifact_name":
    .nemo will be untarred to a temporary folder location and an actual existing path will be returned
else:
    an error will be raised.

警告：在模型的构造函数中使用 .register_artifact 调用。返回的路径不能保证在您退出模型的构造函数后仍然存在。

参数:

  • model – 要为其注册工件的 ModelPT 对象。

  • config_path (str) – 工件键。通常对应于模型配置。

  • src (str) – 工件的路径。

  • verify_src_exists (bool) – 如果设置为 False，则工件是可选的，即使找不到 src，register_artifact 也会返回 None。默认为 True。

返回值:

如果 src 不是 None 或为空，则始终返回绝对路径，保证在模型生命周期内存在

实例生命周期

返回类型:

str

基础 Mixin 类#

class nemo.core.classes.mixins.access_mixins.AccessMixin#

基类： ABC

允许访问模型的中间层输出

register_accessible_tensor(name, tensor)#

注册张量以供后续使用。

classmethod get_module_registry(module: torch.nn.Module)#

从命名的子模块中提取所有注册表，返回字典，其中键是扁平化的模块名称，值是每个此类模块的内部注册表。

reset_registry(registry_key: str | None = None)#

重置所有命名子模块的注册表

property access_cfg#

返回值：跨所有访问 mixin 模块共享的全局访问配置。

class nemo.core.classes.mixins.hf_io_mixin.HuggingFaceFileIO#

基类： ABC

为 NeMo 模型提供 Hugging Face 文件 IO 功能的 Mixin。它通常作为 ModelPT 的 mixin 实现。

此 mixin 提供以下功能： - search_huggingface_models(): 通过一些模型过滤器以编程方式搜索 Hub。 - push_to_hf_hub(): 将模型推送到 Hub。

classmethod get_hf_model_filter() → Dict[str, Any]#

为 HuggingFace 模型生成过滤器。

此外，还包括 Hub 返回结果的一些元数据的默认值。

元数据

resolve_card_info: Bool 标志，如果设置，则返回模型卡元数据。默认值：False。 limit_results: 可选 int，限制返回结果的数量。

返回值:

一个字典，表示可传递给 huggingface list_models() 的参数。

classmethod search_huggingface_models(
model_filter: Dict[str, Any] | None = None,
) → Iterable[huggingface_hub.hf_api.ModelInfo]#

应列出所有可通过 Hugging Face Hub 获得的预训练模型。

以下元数据可以通过 model_filter 传递，以获得更多结果。元数据

resolve_card_info: Bool 标志，如果设置，则返回模型卡元数据。默认值：False。

limit_results: 可选 int，限制返回结果的数量。

# You can replace <DomainSubclass> with any subclass of ModelPT.
from nemo.core import ModelPT

# Get default filter dict
filt = <DomainSubclass>.get_hf_model_filter()

# Make any modifications to the filter as necessary
filt['language'] = [...]
filt['task'] = ...
filt['tags'] = [...]

# Add any metadata to the filter as needed (kwargs to list_models)
filt['limit'] = 5

# Obtain model info
model_infos = <DomainSubclass>.search_huggingface_models(model_filter=filt)

# Browse through cards and select an appropriate one
card = model_infos[0]

# Restore model using `modelId` of the card.
model = ModelPT.from_pretrained(card.modelId)
参数:

model_filter – 可选字典（用于 Hugging Face Hub kwargs），用于过滤返回的兼容模型卡列表，并从每个过滤器中选择所有结果。用户随后可以使用 model_card.modelId 在 from_pretrained() 中恢复 NeMo 模型。

返回值:

ModelInfo 条目的列表。

push_to_hf_hub(
repo_id: str,
*,
pack_nemo_file: bool = True,
model_card: huggingface_hub.ModelCard | None | object | str = None,
commit_message: str = 'Push model using huggingface_hub.',
private: bool = False,
api_endpoint: str | None = None,
token: str | None = None,
branch: str | None = None,
allow_patterns: List[str] | str | None = None,
ignore_patterns: List[str] | str | None = None,
delete_patterns: List[str] | str | None = None,
)#

将模型检查点上传到 Hub。

使用 allow_patterns 和 ignore_patterns 精确过滤应推送到 hub 的文件。使用 delete_patterns 删除同一提交中现有的远程文件。有关更多详细信息，请参见 [upload_folder] 参考。

参数:

  • repo_id (str) – 要推送到的仓库 ID（示例：“username/my-model”）。

  • pack_nemo_file (bool, 可选, 默认为 True) – 是否将模型检查点和配置打包到单个 .nemo 文件中。如果设置为 false，则上传包含模型检查点和配置以及其他工件的目录内容。

  • model_card (ModelCard, 可选) – 要与模型一起上传的模型卡。如果为 None，将使用类本身通过 generate_model_card() 提供的模型卡模板。任何实现 str(obj) 的对象都可以在此处传递。两个关键字替换传递给 generate_model_card()：model_name 和 repo_id。如果模型卡生成字符串，并且它包含 {model_name} 或 {repo_id}，它们将被替换为实际值。

  • commit_message (str, 可选) – 推送时要提交的消息。

  • private (bool, 可选, 默认为 False) – 创建的仓库是否应为私有。

  • api_endpoint (str, 可选) – 将模型推送到 hub 时要使用的 API 端点。

  • token (str, 可选) – 用作远程文件的 HTTP Bearer 授权的令牌。默认情况下，它将使用运行 huggingface-cli login 时缓存的令牌。

  • branch (str, 可选) – 要将模型推送到的 git 分支。默认值为 “main”。

  • allow_patterns (List[str] 或 str, 可选) – 如果提供，则仅推送与至少一个模式匹配的文件。

  • ignore_patterns (List[str] 或 str, 可选) – 如果提供，则不推送与任何模式匹配的文件。

  • delete_patterns (List[str] 或 str, 可选) – 如果提供，则将从仓库中删除与任何模式匹配的远程文件。

返回值:

上传的 HF 仓库的 URL。

神经类型检查#

class nemo.core.classes.common.typecheck(
input_types: TypeState | Dict[str, NeuralType] = TypeState.UNINITIALIZED,
output_types: TypeState | Dict[str, NeuralType] = TypeState.UNINITIALIZED,
ignore_collections: bool = False,
)#

基类： object

一个装饰器，执行输入输出神经类型检查，并将神经类型附加到它包装的函数的输出。

要求类继承自 Typing 才能执行类型检查，否则将引发错误。

# 用法（类级别类型支持）

@typecheck()
def fn(self, arg1, arg2, ...):
    ...

# 用法（函数级别类型支持）

@typecheck(input_types=..., output_types=...)
def fn(self, arg1, arg2, ...):
    ...

需要注意的点

  1. @typecheck() 中的括号 () 是必要的。

    如果没有这些括号，您将遇到 TypeError: __init__() takes 1 positional argument but X were given。

  2. 该函数在定义期间可以接受任意数量的位置参数。

    当您调用此函数时，所有参数都必须仅使用 kwargs 传递。

__call__(wrapped)#

将自身作为函数调用。

class TypeState(value)#

基类：Enum

用于表示提供的类型信息的默认值的占位符。如果此装饰器的构造函数用于覆盖类级别类型定义，则此枚举值表示类型将被覆盖。

wrapped_call(
wrapped,
instance: Typing,
args,
kwargs,
)#

包装器方法，可用于实现 Typing 的类的任何函数。默认情况下，它将使用继承 Typing 的类的 input_types 和 output_types 属性。

可以通过将字典作为参数提供给装饰器来提供局部函数级别的覆盖。

参数:

  • input_types – Union[TypeState, Dict[str, NeuralType]]。默认情况下，使用全局 input_types。

  • output_types – Union[TypeState, Dict[str, NeuralType]]。默认情况下，使用全局 output_types。

  • ignore_collections – Bool。确定是否应断言容器类型的深度检查，或者是否完全跳过深度检查。

static set_typecheck_enabled(enabled: bool = True)#

启用/禁用类型检查的全局方法。

参数:

enabled – bool，当为 True 时将启用类型检查。

static disable_checks()#

上下文管理器，用于在其上下文中临时禁用类型检查。

static set_semantic_check_enabled(enabled: bool = True)#

启用/禁用语义类型检查的全局方法。

参数:

enabled – bool，当为 True 时将启用语义类型检查。

static disable_semantic_checks()#

上下文管理器，用于在其上下文中临时禁用语义类型检查。

神经类型类#

class nemo.core.neural_types.NeuralType(
axes: Any | None = None,
elements_type: Any | None = None,
optional: bool = False,
)#

基类： object

这是表示神经类型概念的主要类。它用于表示输入和输出的类型。

参数:

  • axes (Optional[Tuple]) – AxisTypes 对象的元组，表示每个轴变化的语义。您可以在此处使用简短的、基于字符串的形式。例如：(‘B’, ‘C’, ‘H’, ‘W’) 对应于计算机视觉中常用的 NCHW 格式。(‘B’, ‘T’, ‘D’) 常用于信号处理，表示 [batch, time, dimension/channel]。

  • elements_type (ElementType) – ElementType 类的实例，表示张量内存储内容的语义。例如：logits (LogitsType)、对数概率 (LogprobType) 等。

  • optional (bool) – 默认情况下，此值为 false。如果设置为 True，则表示此类型的端口的输入可以是可选的。

compare(
second,
) → NeuralTypeComparisonResult#

执行 self 与 second 的神经类型比较。当您通过 __call__ 方法链接两个模块的输入/输出时，将调用此比较以确保神经类型兼容性。

compare_and_raise_error(
parent_type_name,
port_name,
second_object,
)#

方法比较一种类型的定义与另一种类型，如果不兼容则引发错误。

class nemo.core.neural_types.axes.AxisType(
kind: AxisKindAbstract,
size: int | None = None,
is_list=False,
)#

基类： object

此类表示轴语义和（可选）其维度 :param kind: 轴的类型是什么？例如 Batch、Height 等 :type kind: AxisKindAbstract :param size: 指定轴是否应具有固定大小。默认情况下，它设置为 None，并且 :type size: int, optional :param 通常不希望为 Batch 和 Time 设置它：:param is_list: 这是否是列表轴或张量轴 :type is_list: bool, default=False

class nemo.core.neural_types.elements.ElementType#

基类： ABC

定义张量元素语义的抽象类。我们依赖 Python 进行继承检查

property type_parameters: Dict[str, Any]#

覆盖此属性以参数化您的类型。例如，您可以使用 ‘dtype’ 关键字指定 ‘storage’ 类型，例如 float、int、bool。另一个示例是，如果您想表示具有特定属性（例如，采样频率）的信号，那么您可以将 sample_freq->value 放在那里。比较两种类型时，它们的 type_parameters 必须匹配。

property fields#

这应该用于逻辑地表示元组/结构。例如，如果您想表示边界框 (x, y, width, height)，您可以在此处放置一个带有名称（‘x’, y’, ‘w’, ‘h’）的元组。在底层，这应该转换为固定大小 = len(fields) 的最后一个张量维度。比较两种类型时，它们的字段必须匹配。

class nemo.core.neural_types.comparison.NeuralTypeComparisonResult(value)#

基类：Enum

比较两个神经类型对象以获得兼容性的结果。当比较 A.compare_to(B) 时

实验管理器#

class nemo.utils.exp_manager.exp_manager(
trainer: lightning.pytorch.Trainer,
cfg: omegaconf.DictConfig | Dict | None = None,
)#

基类

exp_manager 是一个辅助函数，用于管理实验文件夹。它遵循 pytorch lightning 的 exp_dir/model_or_experiment_name/version 范例。如果 lightning trainer 具有 logger，则 exp_manager 将从 logger 获取 exp_dir、name 和 version。否则，它将使用 exp_dir 和 name 参数来创建日志目录。 exp_manager 还允许通过 explicit_log_dir 显式创建文件夹。

版本可以是日期时间字符串或整数。如果 use_datetime_version 设置为 False，则可以禁用日期时间版本。它可选地从 pytorch lightning 创建 TensorBoardLogger、WandBLogger、DLLogger、MLFlowLogger、ClearMLLogger、ModelCheckpoint 对象。它将 sys.argv 和 git 信息（如果可用）复制到日志目录。它为每个进程创建一个日志文件，以将其输出记录到其中。

exp_manager 另外具有恢复功能 (resume_if_exists)，可用于从构造的 log_dir 继续训练。当您需要重复继续训练时（例如在您需要多个连续作业的集群上），您需要避免创建版本文件夹。因此，从 v1.0.0 开始，当 resume_if_exists 设置为 True 时，将忽略创建版本文件夹。

参数:

  • trainer (lightning.pytorch.Trainer) – lightning trainer。

  • cfg (DictConfig, dict) –

    可以具有以下键

    • explicit_log_dir (str, Path): 可用于覆盖 exp_dir/name/version 文件夹创建。默认为

      None，这将使用 exp_dir、name 和 version 来构造日志目录。

    • exp_dir (str, Path): 用于创建日志目录的基本目录。默认为 None，这将日志记录到

      ./nemo_experiments。

    • name (str): 实验的名称。默认为 None，这将通过 name = name 或

      “default” 变为 “default”。

    • version (str): 实验的版本。默认为 None，这将使用日期时间字符串或

      lightning 的 TensorboardLogger 系统来使用 version_{int}。

    • use_datetime_version (bool): 是否使用日期时间字符串作为版本。默认为 True。

    • resume_if_exists (bool): 此实验是否从先前的运行恢复。如果为 True，它将设置

      trainer._checkpoint_connector._ckpt_path，以便 trainer 应该自动恢复。 exp_manager 会将 log_dir 下的文件移动到 log_dir/run_{int}。默认为 False。从 v1.0.0 开始，当 resume_if_exists 为 True 时，我们将不会创建版本文件夹，以便更容易找到下一个运行的日志文件夹。

    • resume_past_end (bool): 如果 resume_if_exists 为 True 并且检查点匹配，则 exp_manager 会出错

      *end.ckpt 指示先前的训练运行已完全完成。此行为可以禁用，在这种情况下，*end.ckpt 将通过将 resume_past_end 设置为 True 来加载。默认为 False。

    • resume_ignore_no_checkpoint (bool): 如果 resume_if_exists 为 True 并且找不到检查点，则 exp_manager 会出错

      可以找到。此行为可以禁用，在这种情况下，exp_manager 将打印一条消息并继续而不恢复，方法是将 resume_ignore_no_checkpoint 设置为 True。默认为 False。

    • resume_from_checkpoint (str): 可用于指定要从中加载的特定检查点文件的路径。这将

      覆盖当 resume_if_exists 为 True 时找到的任何检查点。默认为 None。

    • create_tensorboard_logger (bool): 是否创建 tensorboard logger 并将其附加到 pytorch

      lightning trainer。默认为 True。

    • summary_writer_kwargs (dict): 可以传递给 lightning 的 TensorboardLogger 的 kwargs 字典

      类。请注意，log_dir 由 exp_manager 传递，并且不能存在于此字典中。默认为 None。

    • create_wandb_logger (bool): 是否创建 Weights and Baises logger 并将其附加到 pytorch

      lightning trainer。默认为 False。

    • wandb_logger_kwargs (dict): 可以传递给 lightning 的 WandBLogger 的 kwargs 字典

      类。请注意，如果 create_wandb_logger 为 True，则 name 和 project 是必需参数。默认为 None。

    • create_mlflow_logger (bool): 是否创建 MLFlow logger 并将其附加到 pytorch lightning

      训练。默认为 False

    • mlflow_logger_kwargs (dict): MLFlow logger 的可选参数

    • create_dllogger_logger (bool): 是否创建 DLLogger logger 并将其附加到 pytorch lightning

      训练。默认为 False

    • dllogger_logger_kwargs (dict): DLLogger logger 的可选参数

    • create_clearml_logger (bool): 是否创建 ClearML logger 并将其附加到 pytorch lightning

      训练。默认为 False

    • clearml_logger_kwargs (dict): ClearML logger 的可选参数

    • create_checkpoint_callback (bool): 是否创建 ModelCheckpoint 回调并将其附加到

      pytorch lightning trainer。 ModelCheckpoint 保存 val_loss 最佳的前 3 个模型，*last.ckpt 下的最新检查点，以及训练完成后 *end.ckpt 下的最终检查点。默认为 True。

    • create_early_stopping_callback (bool): 标志，用于决定是否应使用提前停止来停止训练。默认为 False。

      请参阅上面的 EarlyStoppingParams 数据类。

    • create_preemption_callback (bool): 标志，用于决定是否启用抢占回调以保存检查点并在抢占时立即退出训练。

      立即退出训练。默认为 True。

    • create_straggler_detection_callback (bool): 使用落后者检测回调。默认为 False。

    • create_fault_tolerance_callback (bool): 使用容错回调。默认为 False。

    • files_to_copy (list): 要复制到实验日志目录的文件列表。默认为 None，表示

      不复制任何文件。

    • log_local_rank_0_only (bool): 是否仅为本地 rank 0 创建日志文件。默认为 False。

      如果您在使用 DDP 和大量 GPU，并且不希望在 exp 目录中生成大量日志文件，请将其设置为 True。

    • log_global_rank_0_only (bool): 是否仅为全局 rank 0 创建日志文件。默认为 False。

      如果您在使用 DDP 和大量 GPU，并且不希望在 exp 目录中生成大量日志文件，请将其设置为 True。

    • max_time (str): 每次运行的最大挂钟时间。这旨在用于集群，您希望在

      指定时间后保存检查点，并能够从该检查点恢复。默认为 None。

    • seconds_to_sleep (float): 非 rank 0 进程休眠的秒数。用于为 rank 0 初始化提供足够的时间。

    • train_time_interval (timedelta): 传递 timedelta 对象以按 timedelta 保存模型。默认为 None。

      (使用 hydra 的 _target_ 来实现此目的)

返回值:

最终的日志目录，日志文件保存在这里。通常是

exp_dir、name 和 version 的串联。

返回类型:

log_dir (Path)

class nemo.utils.exp_manager.ExpManagerConfig(
explicit_log_dir: str | None = None,
exp_dir: str | None = None,
name: str | None = None,
version: str | None = None,
use_datetime_version: bool | None = True,
resume_if_exists: bool | None = False,
resume_past_end: bool | None = False,
resume_ignore_no_checkpoint: bool | None = False,
resume_from_checkpoint: str | None = None,
create_tensorboard_logger: bool | None = True,
summary_writer_kwargs: ~typing.Dict[~typing.Any,
~typing.Any] | None = None,
create_wandb_logger: bool | None = False,
wandb_logger_kwargs: ~typing.Dict[~typing.Any,
~typing.Any] | None = None,
create_mlflow_logger: bool | None = False,
mlflow_logger_kwargs: ~nemo.utils.loggers.mlflow_logger.MLFlowParams | None = <factory>,
create_dllogger_logger: bool | None = False,
dllogger_logger_kwargs: ~nemo.utils.loggers.dllogger.DLLoggerParams | None = <factory>,
create_clearml_logger: bool | None = False,
clearml_logger_kwargs: ~nemo.utils.loggers.clearml_logger.ClearMLParams | None = <factory>,
create_neptune_logger: bool | None = False,
neptune_logger_kwargs: ~typing.Dict[~typing.Any,
~typing.Any] | None = None,
create_checkpoint_callback: bool | None = True,
checkpoint_callback_params: ~nemo.utils.exp_manager.CallbackParams | None = <factory>,
create_early_stopping_callback: bool | None = False,
early_stopping_callback_params: ~nemo.utils.exp_manager.EarlyStoppingParams | None = <factory>,
create_preemption_callback: bool | None = True,
files_to_copy: ~typing.List[str] | None = None,
log_step_timing: bool | None = True,
log_delta_step_timing: bool | None = False,
step_timing_kwargs: ~nemo.utils.exp_manager.StepTimingParams | None = <factory>,
log_local_rank_0_only: bool | None = False,
log_global_rank_0_only: bool | None = False,
disable_validation_on_resume: bool | None = True,
ema: ~nemo.utils.exp_manager.EMAParams | None = <factory>,
max_time_per_run: str | None = None,
seconds_to_sleep: float = 5,
create_straggler_detection_callback: bool | None = False,
straggler_detection_params: ~nemo.utils.exp_manager.StragglerDetectionParams | None = <factory>,
create_fault_tolerance_callback: bool | None = False,
fault_tolerance: ~nemo.utils.exp_manager.FaultToleranceParams | None = <factory>,
log_tflops_per_sec_per_gpu: bool | None = True,
)#

基类： object

Experiment Manager config for validation of passed arguments.

Exportable#

class nemo.core.classes.exportable.Exportable#

基类： ABC

This Interface should be implemented by particular classes derived from nemo.core.NeuralModule or nemo.core.ModelPT. It gives these entities ability to be exported for deployment to formats such as ONNX.

Usage

# exporting pre-trained model to ONNX file for deployment. model.eval() model.to(‘cuda’) # or to(‘cpu’) if you don’t have GPU

model.export(‘mymodel.onnx’, [options]) # all arguments apart from output are optional.

export(
output: str,
input_example=None,
verbose=False,
do_constant_folding=True,
onnx_opset_version=None,
check_trace: bool | List[torch.Tensor] = False,
dynamic_axes=None,
check_tolerance=0.01,
export_modules_as_functions=False,
keep_initializers_as_inputs=None,
use_dynamo=False,
)#

Exports the model to the specified format. The format is inferred from the file extension of the output file.

参数:

  • output (str) – Output file name. File extension be .onnx, .pt, or .ts, and is used to select export path of the model.

  • input_example (list or dict) – Example input to the model’s forward function. This is used to trace the model and export it to ONNX/TorchScript. If the model takes multiple inputs, then input_example should be a list of input examples. If the model takes named inputs, then input_example should be a dictionary of input examples.

  • verbose (bool) – If True, will print out a detailed description of the model’s export steps, along with the internal trace logs of the export process.

  • do_constant_folding (bool) – If True, will execute constant folding optimization on the model’s graph before exporting. This is ONNX specific.

  • onnx_opset_version (int) – The ONNX opset version to export the model to. If None, will use a reasonable default version.

  • check_trace (bool) – If True, will verify that the model’s output matches the output of the traced model, upto some tolerance.

  • dynamic_axes (dict) – A dictionary mapping input and output names to their dynamic axes. This is used to specify the dynamic axes of the model’s inputs and outputs. If the model takes multiple inputs, then dynamic_axes should be a list of dictionaries. If the model takes named inputs, then dynamic_axes should be a dictionary of dictionaries. If None, will use the dynamic axes of the input_example derived from the NeuralType of the input and output of the model.

  • check_tolerance (float) – The tolerance to use when checking the model’s output against the traced model’s output. This is only used if check_trace is True. Note the high tolerance is used because the traced model is not guaranteed to be 100% accurate.

  • export_modules_as_functions (bool) – If True, will export the model’s submodules as functions. This is ONNX specific.

  • keep_initializers_as_inputs (bool) – If True, will keep the model’s initializers as inputs in the onnx graph. This is ONNX specific.

  • use_dynamo (bool) – If True, use onnx.dynamo_export() instead of onnx.export(). This is ONNX specific.

返回值:

A tuple of two outputs. Item 0 in the output is a list of outputs, the outputs of each subnet exported. Item 1 in the output is a list of string descriptions. The description of each subnet exported can be used for logging purposes.

property disabled_deployment_input_names: List[str]#

Implement this method to return a set of input names disabled for export

property disabled_deployment_output_names: List[str]#

Implement this method to return a set of output names disabled for export

property supported_export_formats: List[ExportFormat]#

Implement this method to return a set of export formats supported. Default is all types.

get_export_subnet(subnet=None)#

Returns Exportable subnet model/module to export

list_export_subnets()#

Returns default set of subnet names exported for this model First goes the one receiving input (input_example)

get_export_config()#

Returns export_config dictionary

set_export_config(args)#

Sets/updates export_config dictionary

On this page