重要提示

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

导出 NeMo 模型#

导出模型#

大多数 NeMo 模型可以导出到 ONNX 或 TorchScript，以便在优化的执行环境（例如 Riva 或 Triton Inference Server）中部署以进行推理。导出接口由 Exportable 混入类提供。如果模型扩展了 Exportable，则可以通过以下方式导出：

from nemo.core.classes import ModelPT, Exportable
# deriving from Exportable
class MyExportableModel(ModelPT, Exportable):
...

mymodel = MyExportableModel.from_pretrained(model_name="MyModelName")
model.eval()
model.to('cuda')  # or to('cpu') if you don't have GPU

# exporting pre-trained model to ONNX file for deployment.
mymodel.export('mymodel.onnx', [options])

如何使用模型导出#

以下参数用于 export()。在大多数情况下，您只需提供输出文件的名称并使用所有默认值

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

output、input_example、verbose、do_constant_folding、onnx_opset_version 选项与 Pytorch onnx.export() 和 jit.trace() 函数中的语义相同，并且会被传递。有关 Pytorch 的 ``onnx.export()`` 的更多信息，请参阅 torch.onnx 函数文档。请注意，如果 input_example 为 None，则会调用 Exportable.input_example()。

output 参数的文件扩展名决定导出格式

  • .onnx->ONNX

  • .pt 或 .ts -> TorchScript。

TorchScript 特定：默认情况下，模块将进行 jit.trace()。您可能需要显式传递一些模块在 jit.script() 下，以便正确跟踪它们。check_trace 参数会传递给 jit.trace()。

ONNX 特定：如果 use_dynamic_axes 为 True，则会使用动态轴调用 onnx.export()。如果 dynamic_axes 为 None，则会从模型的 input_types 定义中推断它们（批次维度是动态的，持续时间等也是如此）。

如果 check_trace 为 True，则生成的 ONNX 也会在 input_example 上运行，并将结果与导出模型的输出进行比较，使用 check_tolerance 参数。请注意较高的默认容差。

如何使模型可导出#

如果您只是使用 NeMo 模型，则前面的示例是您需要了解的全部内容。如果您编写自己的模型，本节重点介绍在扩展 Exportable 后需要注意的事项。

可导出 Hook 和重写#

您通常不需要重写 Exportable 默认方法。但是，Exportable.export() 依赖于您的类中存在某些方法的假设。

@property
def input_example(self) # => Tuple(input, [(input, ...], [Dict])
     """
    Generates input examples for tracing etc.
    Returns:
        A tuple of input examples.
     """

此函数应返回一个元组（通常）为 Tensor - 每个模型输入一个（forward() 的参数）。最后一个元素可以是 Dict，用于按名称指定非位置参数，如 Torch export() 约定。有关更多信息，请参阅 使用字典处理命名参数作为模型输入。

@property
def input_types(self):
@property
def output_types(self):

这些是推断输入/输出名称和动态轴所必需的。如果您的模型派生自 ModulePT，则这些已经存在。另一种常见情况是您的模型包含一个或多个处理输入并生成输出的模块。然后，您应该重写 Exportable 方法 input_module() 和 output_module() 以指向它们，如下例所示

@property
def input_module(self):
    return self.fastpitch

@property
def output_module(self):
    return self.fastpitch

您的模型还应具有导出友好的 forward() 方法 - 这对于 ONNX 和 TorchScript 可能意味着不同的事情。对于 ONNX，您不能有强制命名参数而没有默认值，例如 forward(self, *, text)。对于 TorchScript，您应避免使用 None 并改用 Optional。这些标准非常不稳定，并且可能随每个 PyTorch 版本而更改，因此这是一个反复试验的过程。还有一个普遍的问题，即在许多情况下，用于推理的 forward() 可以简化，甚至可以使用更少的输入/输出。为了解决这个问题，Exportable 在您的模型中查找 forward_for_export() 方法，并使用该方法而不是 forward() 进行导出

# Uses forced named args, many default parameters.
def forward(
    self,
    *,
    text,
    durs=None,
    pitch=None,
    speaker=0,
    pace=1.0,
    spec=None,
    attn_prior=None,
    mel_lens=None,
    input_lens=None,
):
    # Passes through all self.fastpitch outputs
    return self.fastpitch(
        text=text,
        durs=durs,
        pitch=pitch,
        speaker=speaker,
        pace=pace,
        spec=spec,
        attn_prior=attn_prior,
        mel_lens=mel_lens,
        input_lens=input_lens,
    )


# Uses less inputs, no '*', returns less outputs:
def forward_for_export(self, text):
    (
        spect,
        durs_predicted,
        log_durs_predicted,
        pitch_predicted,
        attn_soft,
        attn_logprob,
        attn_hard,
        attn_hard_dur,
        pitch,
    ) = self.fastpitch(text=text)
    return spect, durs_predicted, log_durs_predicted, pitch_predicted

为了与 input_types()/output_types() 保持一致，Exportable 中也存在这些 hook，可让您从导出过程中排除特定的输入/输出

@property
def disabled_deployment_input_names(self):
    """Implement this method to return a set of input names disabled for export"""
    return set(["durs", "pitch", "speaker", "pace", "spec", "attn_prior", "mel_lens", "input_lens"])

@property
def disabled_deployment_output_names(self):

导出模型的另一个常见要求是在导出之前运行某些网络修改以提高推理效率 - 例如禁用某些卷积中的掩码或删除批归一化。更好的风格是在 ModelPT.eval() 上进行这些操作（并在 .train() 上反转），但这并非总是可行，因此 Exportable 中提供了以下 hook 来运行这些操作

def _prepare_for_export(self, **kwargs):
    """
    Override this method to prepare module for export. This is in-place operation.
    Base version does common necessary module replacements (Apex etc)
    """
# do graph modifications specific for this model
    replace_1D_2D = kwargs.get('replace_1D_2D', False)
    replace_for_export(self, replace_1D_2D)
# call base method for common set of modifications
    Exportable._prepare_for_export(self, **kwargs)

某些需要控制流的模型需要分多个部分导出。RNNT 网络是典型的例子。为了方便这一点，提供了以下 hook。例如，要导出模型的“encoder”和“decoder”子网，请重载 list_export_subnets 以返回 ['encoder', 'decoder']。

def get_export_subnet(self, subnet=None):
    """
    Returns Exportable subnet model/module to export
    """


def list_export_subnets(self):
    """
    Returns default set of subnet names exported for this model
    First goes the one receiving input (input_example)
    """

某些网络可能会根据用户可设置的选项（例如用于 TTS 的 Ragged Batch 支持或用于 ASR 的缓存支持）以不同的方式导出。为了方便这一点，Exportable 提供了 set_export_config() 方法，用于将键/值对设置到预定义的 model.export_config 字典，以在导出期间使用

def set_export_config(self, args):
    """
    Sets/updates export_config dictionary
    """

此外，如果需要设置配置的操作 hook，Exportable 后代可以重载此方法以包含一个。可以在 <NeMo_git_root>/nemo/collections/asr/models/rnnt_models.py 中找到示例。

以下是在 <NeMo_git_root>/scripts/export.py 中，set_export_config() 调用如何与命令行参数绑定的示例

python scripts/export.py  hybrid_conformer.nemo hybrid_conformer.onnx --export-config decoder_type=ctc

可导出的模型代码#

最重要的是，模型中的实际 Torch 代码应与 ONNX 或 TorchScript 兼容（理想情况下，两者都兼容）。#. 确保代码以 Torch 方式编写 - 避免使用裸 Numpy 或 Python 操作数。#. 创建您的模型 Exportable 并添加一个导出单元测试，以立即捕获 ONNX/TorchScript 中不支持的任何操作/构造。

有关更多信息，请参阅 PyTorch 文档

在本页中