如何在 Riva 上部署使用 NeMo 训练的自定义声学模型 (Conformer-CTC)#

本教程将引导您完成在 NVIDIA Riva 上部署使用 NVIDIA NeMo 训练的自定义声学模型 (Conformer-CTC) 的过程。

NVIDIA Riva 概述#

NVIDIA Riva 是一个 GPU 加速的 SDK,用于构建为您的用例定制并提供实时性能的语音 AI 应用程序。
Riva 提供丰富的语音和自然语言理解服务,例如

  • 自动语音识别 (ASR)。

  • 文本到语音合成 (TTS)。

  • 一系列自然语言处理 (NLP) 服务,例如命名实体识别 (NER)、标点符号和意图分类。

在本教程中,我们将部署一个使用 NeMo 在 Riva 上训练的自定义声学模型 (Conformer-CTC)。
要了解 Riva ASR API 的基础知识,请参考 Riva ASR Python 入门

有关 Riva 的更多信息,请参阅 Riva 产品页面Riva 开发者文档

NeMo (神经模块) 和 nemo2riva#

NVIDIA NeMo 是一个开源框架,用于使用简单的 Python 接口构建、训练和微调 GPU 加速的语音 AI 和自然语言理解 (NLU) 模型。要使用 NeMo 微调 Conformer-CTC 声学模型,请参阅 Conformer-CTC 微调教程

nemo2riva 命令行工具提供了导出 .nemo 模型的功能,该模型可以部署在使用 NVIDIA Riva 的格式中,NVIDIA Riva 是一个高性能应用程序框架,用于使用 GPU 的多模态对话式 AI 服务。nemo2riva 的 Python .whl 文件包含在 Riva Quick Start 资源文件夹中。您也可以使用 pip 安装 nemo2riva,如 Conformer-CTC 微调教程 中所示。

本教程探讨如何使用 .riva 模型(调用 nemo2riva CLI 工具的结果,请参阅 Conformer-CTC 微调教程)并利用 Riva ServiceMaker 框架来聚合 Riva 部署到目标环境所需的所有必要工件。一旦模型部署在 Riva 中,您就可以向服务器发出推理请求。我们将演示整个过程是多么快速和直接。在本教程中,您将学习如何

  • 使用 Riva ServiceMaker 从 .riva 文件构建 .rmir 模型管道。

  • 在本地 Riva 服务器上部署模型。

  • 使用 Riva API 绑定从演示客户端发送推理请求。


前提条件#

在我们开始之前,请确保您已拥有

  • 访问 NVIDIA NGC 并能够下载 Riva Quick Start 资源

  • 您要部署的 .riva 模型文件。您可以使用 nemo2riva 命令将 .nemo 模型文件转换为 .riva 模型文件。

    • 有关使用 NeMo 自定义 Conformer-CTC 声学模型并使用 nemo2riva 导出结果模型的更多信息,请参阅 Conformer-CTC 微调教程

    • 或者,您可以从此处获取英语 ASR 的预训练 Conformer-CTC .riva 模型。

    • 有关训练 NeMo 模型的更多信息,请参阅 NeMo 文档中的“训练”部分。

    • 有关 Conformer-CTC 架构的更多信息,请参阅 NeMo ASR 模型页面中的“Conformer-CTC”部分。

    • 有关使用 NeMo 训练 Conformer-CTC 所需的配置文件,请参阅 NeMo ASR 模型配置文件页面中的“Conformer-CTC”部分。


Riva ServiceMaker#

Riva ServiceMaker 是一组工具,用于聚合 Riva 部署到目标环境所需的所有必要工件(模型、文件、配置和用户设置)。它有两个主要组件

Riva-Build#

此步骤有助于构建 Riva 就绪的模型版本。其唯一输出是 Riva 中支持的服务的端到端管道的中间格式(称为 RMIR)。让我们考虑一个 Conformer-CTC ASR 模型。

riva-build 负责将一个或多个导出的模型(.riva 文件)组合成一个包含中间格式的文件,该中间格式称为 Riva 模型中间表示 (RMIR)(.rmir)。此文件包含整个端到端管道的部署不可知规范,以及最终部署和推理所需的所有资产。有关更多信息,请参阅文档

# IMPORTANT: UPDATE THESE PATHS 

# Riva Docker
RIVA_CONTAINER = "<add container name>"

# Example: 
# RIVA_CONTAINER = f"nvcr.io/nvidia/riva/riva-speech:{__riva_version__}"

# Directory where the .riva model is stored $MODEL_LOC/*.riva
MODEL_LOC = "<add path to model location>"

# Name of the .riva file
MODEL_NAME = "<add model name>"
# Get the Riva Docker container
! docker pull $RIVA_CONTAINER

如果子目录尚不存在,请在 MODEL_LOC 中创建一个子目录来存储您的 .rmir 文件。

! mkdir -p $MODEL_LOC/rmir

构建 .rmir 文件。#

注意

  1. 如果您在调用 nemo2riva 时通过添加 --key 标志加密了您的声学模型和/或语言模型,或者您下载了 2023 年之前 NGC 发布的预训练模型,则需要在 riva-build 命令中模型的名称后附加一个冒号,然后附加密钥的值,如下所示。您可能会发现设置名为 KEY 的字符串变量并将其作为 $KEY 传递到相应的 riva-build 参数中很方便。旧版预训练模型的标准加密密钥是 tlt_encode

  2. 如果您想为美国英语以外的受支持语言构建 ASR 管道,请参阅 Riva ASR 管道配置文档。要获取特定应用程序的正确 riva-build 参数,请从页面第一部分底部的交互式 Web 菜单中选择声学模型、语言和管道类型(本教程的目的为离线)。

# Syntax: riva-build <task-name> output-dir-for-rmir/model.rmir:key dir-for-riva/model.riva:key
! docker run --rm --gpus 0 -v $MODEL_LOC:/data $RIVA_CONTAINER -- \
    riva-build speech_recognition \
        /data/rmir/asr_offline_conformer_ctc.rmir:$KEY \
        /data/$MODEL_NAME:$KEY \
        --offline \
        --name=asr_offline_conformer_ctc_pipeline \
        --decoder_type=greedy \
        --ms_per_timestep=40 \
        --chunk_size=4.8 \
        --left_padding_size=1.6 \
        --right_padding_size=1.6 \
        --max_batch_size=16 \
        --nn.fp16_needs_obey_precision_pass \
        --featurizer.use_utterance_norm_params=False \
        --featurizer.precalc_norm_time_steps=0 \
        --featurizer.precalc_norm_params=False \
        --featurizer.max_batch_size=512 \
        --featurizer.max_execution_batch_size=512 \
        --language_code=en-US

Riva-Deploy#

部署工具将一个或多个 RMIR 文件和一个目标模型存储库目录作为输入。它创建一个集合配置,指定执行管道,最后将所有这些资产写入输出模型存储库目录。

注意: 如果您在构建 .rmir 文件时使用 riva-build 向其添加了加密密钥,请确保在 riva-deploy 命令中模型的名称后附加一个冒号,然后附加密钥的值,如下所示。

# Syntax: riva-deploy -f dir-for-rmir/model.rmir:key output-dir-for-repository
! docker run --rm --gpus 0 -v $MODEL_LOC:/data $RIVA_CONTAINER -- \
    riva-deploy -f  \
        /data/rmir/asr_offline_conformer_ctc.rmir:$KEY \
        /data/models/

启动 Riva 服务器#

生成模型存储库后,我们就可以启动 Riva 服务器了。首先,从 NGC 下载 Riva Quick Start 资源。在此处设置目录路径

# Set the Riva Quick Start directory
RIVA_DIR = "<Path to the uncompressed folder downloaded from quickstart(include the folder name)>"

接下来,我们修改 config.sh 文件以启用相关的 Riva 服务(Conformer-CTC 模型的 ASR),提供加密密钥,以及上一步中生成的模型存储库路径 (riva_model_loc) 以及其他配置。

例如,如果在上面的 $MODEL_LOC/models 中生成模型存储库,则可以将 riva_model_loc 指定为与 MODEL_LOC 相同的目录。

models_asr/nlp/tts 中指定的模型的预训练版本是从 NGC 获取的。由于我们使用的是自定义模型,因此我们可以在 models_asr 中注释掉它(以及与您的用例无关的任何其他模型)。

config.sh 代码片段#

# Enable or Disable Riva Services
service_enabled_asr=true 
service_enabled_nlp=true # MAKE CHANGES HERE - SET TO FALSE
service_enabled_tts=true # MAKE CHANGES HERE - SET TO FALSE
service_enabled_nmt=true # MAKE CHANGES HERE - SET TO FALSE

...

# Locations to use for storing models artifacts
#
# If an absolute path is specified, the data will be written to that location
# Otherwise, a Docker volume will be used (default).
#
# riva_init.sh will create a `rmir` and `models` directory in the volume or
# path specified.
#
# RMIR ($riva_model_loc/rmir)
# Riva uses an intermediate representation (RMIR) for models
# that are ready to deploy but not yet fully optimized for deployment. Pretrained
# versions can be obtained from NGC (by specifying NGC models below) and will be
# downloaded to $riva_model_loc/rmir by `riva_init.sh`
#
# Custom models produced by NeMo or TLT and prepared using riva-build
# may also be copied manually to this location $(riva_model_loc/rmir).
#
# Models ($riva_model_loc/models)
# During the riva_init process, the RMIR files in $riva_model_loc/rmir
# are inspected and optimized for deployment. The optimized versions are
# stored in $riva_model_loc/models. The riva server exclusively uses these
# optimized versions.
riva_model_loc="riva-model-repo"  # MAKE CHANGES HERE - REPLACE WITH $MODEL_LOC

if [[ $riva_target_gpu_family == "tegra" ]]; then
    riva_model_loc="`pwd`/model_repository"
fi

# The default RMIRs are downloaded from NGC by default in the above $riva_rmir_loc directory
# If you'd like to skip the download from NGC and use the existing RMIRs in the $riva_rmir_loc
# then set the below $use_existing_rmirs flag to true. You can also deploy your set of custom
# RMIRs by keeping them in the riva_rmir_loc dir and use this quickstart script with the
# below flag to deploy them all together.
use_existing_rmirs=false        # MAKE CHANGES HERE - set to true                    

注意 在继续之前,请确保执行以下操作

或者手动执行这些任务

  1. 在 Jupyter Lab 的文件导航器中,导航到 $RIVA_DIR 并打开 config.sh

  2. 按照上面的代码片段配置设置

    • 将 NLP、TTS 和 NMT 服务设置为 false

    • riva_model_loc 路径设置为也分配给 MODEL_LOC 的路径

    • 将变量 use_existing_rmirs 设置为 true

或者运行下面的单元格

ENABLE_ASR = 'true'
ENABLE_NLP = 'false'
ENABLE_TTS = 'false'
ENABLE_NMT = 'false'

!sed -i "s|service_enabled_asr=.*|service_enabled_asr=$ENABLE_ASR|g" $RIVA_DIR/config.sh
!sed -i "s|service_enabled_nlp=.*|service_enabled_nlp=$ENABLE_NLP|g" $RIVA_DIR/config.sh
!sed -i "s|service_enabled_tts=.*|service_enabled_tts=$ENABLE_TTS|g" $RIVA_DIR/config.sh
!sed -i "s|service_enabled_nmt=.*|service_enabled_nmt=$ENABLE_NMT|g" $RIVA_DIR/config.sh

!sed -i "/\sriva_model_loc=.*/! s|riva_model_loc=.*|riva_model_loc=\"$MODEL_LOC\"|g" $RIVA_DIR/config.sh

!sed -i "s|use_existing_rmirs=.*|use_existing_rmirs=true|g" $RIVA_DIR/config.sh
# Ensure you have permission to execute these scripts
! cd $RIVA_DIR && chmod +x ./riva_init.sh && chmod +x ./riva_start.sh
# Run Riva Init. This will fetch the containers/models
# YOU CAN SKIP THIS STEP IF YOU DID RIVA DEPLOY
! cd $RIVA_DIR && ./riva_init.sh config.sh
# Run Riva Start. This will deploy your model(s).
! cd $RIVA_DIR && ./riva_start.sh config.sh

运行推理#

在 Riva 服务器启动并运行您的模型后,您可以发送推理请求来查询服务器。

要发送 gRPC 请求,您可以为客户端安装 Riva Python API 绑定。这可以作为 PyPI 上的 Python 模块使用。

# Install the Client API Bindings
! pip install nvidia-riva-client
import riva.client

连接到 Riva 服务器并运行推理#

调用此推理函数会查询 Riva 服务器(使用 gRPC)以转录音频文件。

def run_inference(audio_file, server='localhost:50051', print_full_response=False):
    with open(audio_file, 'rb') as fh:
        data = fh.read()
    
    auth = riva.client.Auth(uri=server)
    client = riva.client.ASRService(auth)
    config = riva.client.RecognitionConfig(
        language_code="en-US",
        max_alternatives=1,
        enable_automatic_punctuation=False,
    )
    
    response = client.offline_recognize(data, config)
    if print_full_response: 
        print(response)
    else:
        print(response.results[0].alternatives[0].transcript)

现在我们可以实际查询 Riva 服务器了。

audio_file = "<add path to .wav (PCM-, A-Law-, or U-Law-encoded), .flac, .opus, or .ogg (Opus-encoded) file>"
run_inference(audio_file)

您可以在关闭 Jupyter 内核之前停止 Riva 服务器。

! cd $RIVA_DIR && ./riva_stop.sh