深度学习模型部署(1):模型导出

在本系列文章的这一部分，我们将从模型训练与优化，转向模型在实际应用中的落地：学习如何导出训练好的模型，使用 Triton 进行部署，评估并优化推理性能，最后构建一个可以调用推理服务的简单应用，完成从开发到生产的闭环流程。

我们将以 NVIDIA 的Triton Inference Server为例，它可以帮助我们更高效地管理和提供推理服务，支持多种部署场景，包括本地服务器、数据中心和云端平台。在这篇文章中，我们先来学习如何将使用PyTorch训练的BERT模型检查点导出到NVIDIA Triton推理服务器。

本文用到的脚本文件可在Triton模型部署相关脚本文件下载。

1 NVIDIA Triton Inference Server

这张图展示了NVIDIA Triton Inference Server(英伟达Triton推理服务器)的整体架构与工作流程：

(1)Client Application(客户端应用)

最上层是客户端应用程序，这可以是任意终端，如笔记本、手机或平板，运行你的推理请求，比如一个文本分类服务、聊天机器人或搜索系统。

客户端通过以下方式与 Triton 通信：

• Python/C++ Client Library：官方提供的客户端 SDK，用于封装请求/响应逻辑。
• 支持两种通信协议：
  • HTTP：适合 Web 接口、REST API 场景。
  • gRPC：适合高性能场景，支持多路复用和更低延迟。
• 也可以直接通过 C API 接口进行调用(适用于低级系统集成)。

(2)Model Repository(模型仓库)

模型文件被放置在一个持久化的模型仓库中(通常是一个挂载的卷或远程存储)，Triton 会自动从这个目录中加载模型。模型的配置、版本和优化信息都来源于这里，可以支持热更新(无中断更新模型)。

(3)NVIDIA Triton Inference Server 核心组件

这部分是 Triton 的核心部分，负责调度、执行推理，并返回结果。

1. Model Management(模型管理)
   Triton 会从模型仓库加载模型，并管理其生命周期。你可以指定多个模型同时部署，支持版本控制。
2. Inference Request / Response(推理请求与响应)
   请求进入服务器后被解析，分发给对应模型的执行队列，然后进行推理处理，最后返回响应。
3. Per-model Scheduler Queues(模型调度队列)
   每个模型有自己的调度队列，用于支持动态批处理(Dynamic Batching)、请求排序等调度策略，提高吞吐量。
4. Framework Backends(框架后端)
   Triton 支持多种深度学习框架的模型推理，图中列出了主要支持的格式和框架：TensorRT、TensorFlow、ONNX、PyTorch和自定义后端(Custom)
5. Scheduler(调度器)
   调度器根据资源利用率、请求优先级等因素分发请求到后端推理引擎。
6. GPU/CPU 执行资源
   支持多 GPU 并行推理，也可以使用 CPU。你可以在配置中指定每个模型使用哪些硬件资源。

(4)Status/Health Metrics Export(状态与健康监控)

Triton 会通过 HTTP 接口暴露一系列监控指标(如吞吐量、延迟、模型加载状态等)，用于系统监控与性能评估。这对于生产环境中的稳定性和可运维性至关重要。

2 优化与性能

对训练好的模型进行优化通常能够显著提升推理性能，主要体现在提高带宽、降低延迟等方面。即使项目不采用知识蒸馏或模型剪枝等高级优化手段，仅仅借助常规的模型优化工具，也可以实现明显的性能改善。下图展示了同一模型在三种部署方式下的推理效果差异：原始的 TensorFlow 模型、经过 TensorRT 后处理的模型，以及使用 TensorRT 进行完整优化的模型。

现代推理服务器通常支持多种模型格式，以满足不同项目、工具和用户偏好的需求。由于之前的基于Tranformer的NLP实战中我们使用PyTorch训练了一个BERT检查点，并计划将其部署到Triton推理服务器上，因此我们将专注于部署基于PyTorch的模型。

虽然在 PyTorch 中训练了模型，但部署到Triton时，可以通过将模型转换为TorchScript、ONNX或直接生成TensorRT引擎等方式，以不同的格式和后端执行方式运行这些模型。

1. PyTorch JIT / TorchScript

   TorchScript 是 PyTorch 官方支持的模型导出格式，可将动态图转换为可序列化、可部署的静态图。它适合直接将 PyTorch 模型导出为部署版本，无需转换框架。Triton 可以通过 libtorch 后端加载和执行 TorchScript 模型，是最接近 PyTorch 原生方式的部署方案，部署过程简单，调试也较为直观。

2. ONNX Runtime

   ONNX(Open Neural Network Exchange)是一个开放的模型交换格式，适用于在不同框架间共享模型。PyTorch 模型可以被导出为 ONNX 格式，并通过 ONNX Runtime 高效推理。此方式适合需要跨平台、跨框架兼容性的部署场景，同时支持在 CPU 或 GPU 上运行，易于移植。

3. ONNX-TensorRT

   在导出为 ONNX 格式后，可以进一步利用 TensorRT 对模型进行图优化、算子融合和精度压缩(如 FP16、INT8)，生成高效的TensorRT引擎。这种方式结合了ONNX的兼容性和TensorRT的高性能，适用于需要高吞吐、低延迟的 GPU 推理场景，是生产部署中常见的优化路径。

4. TensorRT

   TensorRT是NVIDIA推出的专用推理引擎，支持直接加载经过转换的 TensorRT 引擎文件(如从 PyTorch 通过中间工具转化)。它提供了最极致的性能优化，包括动态批处理、精度量化等，但转换过程相对复杂，适用于对性能要求极高、可控性强的部署环境。

虽然我们当前用的是 PyTorch + BERT 模型，但Triton还能支持很多别的格式和框架，因此在一个统一的平台下可以部署多个来源的模型。

1. TensorFlow GraphDef
   TensorFlow 的早期模型格式，是用 .pb 文件保存的静态图(graph definition)。如果你使用 TensorFlow 1.x 训练模型，可以直接导出为 GraphDef，然后在 Triton 中部署。
2. TensorFlow SavedModel
   TensorFlow 2.x推荐使用的模型格式，支持保存计算图和权重。SavedModel更完整、可扩展，适合现代 TensorFlow 项目，也是 Triton 原生支持的一种格式。
3. Caffe2 导出格式
   虽然Caffe2现在较少使用(已被PyTorch整合)，但Triton仍然支持其模型部署。适用于那些历史上使用 Caffe2 构建的工业项目。
4. 自定义模型(Custom Backend)
   如果你的模型格式不在上述列表中，或需要特殊的前后处理逻辑，可以自己写一个“自定义后端”作为Triton插件。这通常是一个可执行文件或动态链接库(如 .so)，Triton会调用它来完成模型推理。

在本节中，我们将探索上述部分部署方式，并观察它们对性能的影响。同时也会尝试关键配置参数，例如批大小和数值精度(FP32 和 FP16)。

3 导出BERT检查点

我们要部署的 BERT 模型检查点文件bert_qa.pt应保存在data目录下。bert_qa.pt是一个使用 BERT-Large 结构、已经训练好、并针对SQuAD问答任务微调过的模型，它可以直接用于问答类推理任务。

辅助脚本

在探索各种部署配置时，我们将重复一些步骤。为了简化流程并专注于配置和结果，我们将使用一些辅助脚本自动化部分操作。您可以自行查看这些脚本的源码：

• utilities/wait_for_triton_server.sh：通过API检查Triton服务器是否已在线和就绪
• deployer/deployer.py：将检查点转换为可部署模型并导出
• utilities/run_perf_client_local.sh：使用perf_client工具进行性能测试

Triton 服务器已经在容器中部署完毕，并通过主机名 “triton” 和端口 8000 提供服务。运行以下单元，检查服务器是否返回 “200 OK”：

# 设置 Triton 服务器主机名并检查连接状态
tritonServerHostName = "triton"
!./utilities/wait_for_triton_server.sh {tritonServerHostName}

3.1 Triton 模型仓库

启动Triton Server时，通常会配置其观察一个本地或远程文件系统中的模型目录。该目录称为模型仓库。用于启动Triton Server的命令中会指定该路径，例如：

# 启动 Triton 时指定模型仓库路径
tritonserver --model-repository="/path/to/model/repository"

模型仓库需要满足以下目录结构：

<model-repository-path>/
  <model-name>/
    [config.pbtxt]
    [<output-labels-file> ...]
    <version>/
      <model-definition-file>
    <version>/
      <model-definition-file>
  ...

本实验所用容器默认配置使用 ./model_repository 文件夹作为模型仓库，因此该文件夹内的任何更改都会影响 Triton Server 的行为。

向 Triton 注册一个新模型需执行以下步骤：

1. 在模型仓库中创建一个新的模型文件夹。该文件夹名称应与希望对外暴露的服务名称一致。
2. 在模型文件夹中创建config.pbtxt文件，提供基本的服务配置。
3. 在模型文件夹中创建至少一个模型版本的子文件夹。子文件夹名称为模型版本号，您可以托管多个版本的同一模型。

接下来，我们将演示如何将模型导出到Triton。

3.2 TorchScript 导出

现在我们要完成：

• 将 PyTorch 检查点转换为TorchScript
• 生成Triton所需的配置文件
• 将生成的模型文件部署到模型仓库中

请运行以下代码单元。由于要加载模型并执行转换，过程可能会持续一分钟左右。

# 设置模型名称
modelName = "bertQA-torchscript"
# 使用 deployer.py 导出模型到 TorchScript，并配置 Triton 所需信息
!python ./deployer/deployer.py \
    --ts-script \
    --save-dir ./candidatemodels \
    --triton-model-name {modelName} \
    --triton-model-version 1 \
    --triton-max-batch-size 8 \
    --triton-dyn-batching-delay 0 \
    --triton-engine-count 1 \
    -- --checkpoint "/dli/task/data/bert_qa.pt" \
    --config_file ./bert_config.json \
    --vocab_file ./vocab \
    --predict_file ./squad/v1.1/dev-v1.1.json \
    --do_lower_case \
    --batch_size=8 

该脚本加载 bert_qa.pt 检查点，以 TorchScript 格式导出模型到 bertQA-torchscript 文件夹，并标记为版本 1。

现在我们在./candidatemodels/bertQA-torchscript/中可以看到模型以model.pt文件保存为TorchScript格式，同时生成了配置文件 config.pbtxt。配置文件定义了：

• 模型名称
• 使用的平台类型：此处为 pytorch_libtorch
• 网络的输入输出维度
• 使用的优化方式（默认 GPU 上的 TorchScript 优化）
• 实例组设置：此处为一个模型实例驻留在 GPU 0 上

接下来将模型文件夹移动至 Triton 的模型仓库：

!mv ./candidatemodels/bertQA-torchscript model_repository/

3.3 测试导出结果

运行以下代码单元以启动推理过程，并对推理性能进行简单测量。首先，我们将设置一些测试参数。这里将 maxConcurrency 设置为 2，意味着压力测试将运行两次：第一次只使用一个线程，第二次使用两个线程同时向服务器发送请求。在未开启模型并发执行或动态批处理功能的情况下，你认为这种设置对性能会有何影响？

• 带宽会增加还是减少？
• 延迟会增加还是减少？

# 设置测试参数
modelVersion="1"
precision="fp32"
batchSize="1"
maxLatency="500"
maxClientThreads="10"
maxConcurrency="2"
dockerBridge="host"
resultsFolderName="1"
profilingData="utilities/profiling_data_int64"
# 运行本地性能测试脚本
!./utilities/run_perf_client_local.sh \
                    {modelName} \
                    {modelVersion} \
                    {precision} \
                    {batchSize} \
                    {maxLatency} \
                    {maxClientThreads} \
                    {maxConcurrency} \
                    {tritonServerHostName} \
                    {dockerBridge} \
                    {resultsFolderName} \
                    {profilingData}

如果一切正常，您将看到类似下图的输出结果，展示了不同并发配置下的推理性能：

如果遇到error: failed to get model metadata错误，请尝试重新运行上面的代码。

3.4 使用ONNX格式部署模型

接下来我们将尝试另一种模型部署方式，即ONNX(Open Neural Network Exchange)。ONNX是一种开放格式，专用于表示和交换神经网络模型。它定义了一组常用算子，以及用于模型交换的标准文件格式。ONNX的优势在于它被广泛支持，能够在多种深度学习工具中进行转换或部署，包括TensorRT。

与之前类似，我们将导出模型，但这次使用ONNX格式。我们仍然使用前面用过的导出工具，只需将导出格式从 ts-script更改为onnx：

# 设置模型名称和导出格式
modelName = "bertQA-onnx"
exportFormat = "onnx"
# 使用 deployer.py 脚本导出为 ONNX 格式模型
!python ./deployer/deployer.py \
    --{exportFormat} \
    --save-dir ./candidatemodels \
    --triton-model-name {modelName} \
    --triton-model-version 1 \
    --triton-max-batch-size 8 \
    --triton-dyn-batching-delay 0 \
    --triton-engine-count 1 \
    -- --checkpoint ./data/bert_qa.pt \
    --config_file ./bert_config.json \
    --vocab_file ./vocab \
    --predict_file ./squad/v1.1/dev-v1.1.json \
    --do_lower_case \
    --batch_size=8

与TorchScript 序列化格式类似，ONNX 格式也可以轻松查看和理解(部分内容是可读的)。我们可以在./candidatemodels/bertQA-onnx/1查看配置文件和模型文件，模型以ONNX格式保存。

我们有两个选择来在Triton中执行该ONNX模型：

• 使用ONNX RuntimE运行模型
• 使用TensorRT将ONNX转换为TensorRT引擎执行

我们将尝试这两种方法，并比较它们对推理性能的影响。现在先将 ONNX 模型移动到模型仓库中：

# 部署 ONNX 模型到 Triton 模型仓库
!mv ./candidatemodels/bertQA-onnx model_repository/

然后在10个并发级别下运行压力测试：

# 配置参数并运行性能测试脚本
modelName = "bertQA-onnx"
maxConcurrency = "10"
batchSize = "8"
print("Running: "+modelName)
!bash ./utilities/run_perf_client_local.sh \
                    {modelName} \
                    {modelVersion} \
                    {precision} \
                    {batchSize} \
                    {maxLatency} \
                    {maxClientThreads} \
                    {maxConcurrency} \
                    {tritonServerHostName} \
                    {dockerBridge} \
                    {resultsFolderName} \
                    {profilingData}

运行完后可以观察结果：我们是否成功在所有10个并发级别下完成了基准测试？(是否因超时提前结束？)，请求延迟是否超出了我们配置的500毫秒上限？

接下来我们重新导出一次ONNX模型，以便配置TensorRT执行：

# 重新设置模型名称用于 TensorRT 配置
modelName = "bertQA-onnx-trt-fp16"
exportFormat = "onnx"
# 以 ONNX 格式导出模型用于后续 TensorRT 优化
!python ./deployer/deployer.py \
    --{exportFormat} \
    --save-dir ./candidatemodels \
    --triton-model-name {modelName} \
    --triton-model-version 1 \
    --triton-max-batch-size 8 \
    --triton-dyn-batching-delay 0 \
    --triton-engine-count 1 \
    -- --checkpoint ./data/bert_qa.pt \
    --config_file ./bert_config.json \
    --vocab_file ./vocab \
    --predict_file ./squad/v1.1/dev-v1.1.json \
    --do_lower_case \
    --batch_size=8

这条命令将在./candidatemodels/bertQA-onnx-trt-fp16/中再次生成ONNX模型及配置文件：

3.5 启用 TensorRT 优化

为了启用TensorRT，我们需要在config.pbtxt文件中添加以下配置段：

optimization {
   execution_accelerators {
      gpu_execution_accelerator : [ {
         name : "tensorrt"
         parameters { key: "precision_mode" value: "FP16" }
      }]
   }
cuda { graphs: 0 }
}

添加的config.pbtxt文件如下：

name: "bertQA-onnx-trt-fp16"
platform: "onnxruntime_onnx"
max_batch_size: 8
input [
{
    name: "input__0"
    data_type: TYPE_INT64
    dims: [384]
},
{
    name: "input__1"
    data_type: TYPE_INT64
    dims: [384]
},
{
    name: "input__2"
    data_type: TYPE_INT64
    dims: [384]
}
]
output [
{
    name: "output__0"
    data_type: TYPE_FP32
    dims: [384]
}, 
{
    name: "output__1"
    data_type: TYPE_FP32
    dims: [384]
}
]
optimization {
    execution_accelerators {
      gpu_execution_accelerator : [ {
         name : "tensorrt"
         parameters { key: "precision_mode" value: "FP16" }
      }]
   }
  cuda {
    graphs: 0
  }
}
instance_group [
    {
        count: 1
        kind: KIND_GPU
        gpus: [ 0 ]
    }
]

然后运行以下命令将模型移动到模型仓库中：

# 将优化后的模型移动到模型仓库
!mv ./candidatemodels/bertQA-onnx-trt-fp16 model_repository/

运行性能测试工具以分析优化效果。注意：首次加载模型到TensorRT引擎可能需要一些时间。

# 运行性能测试脚本以验证 TensorRT 优化效果
modelName = "bertQA-onnx-trt-fp16"
maxConcurrency= "10"
batchSize="8"
print("Running: " + modelName)
!bash ./utilities/run_perf_client_local.sh \
                    {modelName} \
                    {modelVersion} \
                    {precision} \
                    {batchSize} \
                    {maxLatency} \
                    {maxClientThreads} \
                    {maxConcurrency} \
                    {tritonServerHostName} \
                    {dockerBridge} \
                    {resultsFolderName} \
                    {profilingData}

最后，让我们比较ONNX runtime和TensorRT的性能：

(1)延迟是如何变化的，尤其是在高并发下？

ONNX Runtime在高并发下延迟明显上升，甚至超出设定阈值导致超时；而TensorRT延迟控制良好，即使在并发为10时也保持稳定，展现出更强的性能和可扩展性。

(2)带宽有何变化？是否能解释观察到的带宽差异？

TensorRT 在高并发下仍能持续提升带宽，充分利用 GPU 资源；而ONNX Runtime的带宽提升趋于饱和，说明其缺乏深度硬件优化，限制了吞吐能力。

(3)为什么ONNX模型在小于10的并发下会超时？而TensorRT模型在并发为10时的延迟又是怎样的？

ONNX 模型由于未充分优化，在多请求排队时易出现瓶颈，导致超时；而TensorRT使用了图优化与FP16加速，即使在并发10时也能保持在延迟阈值内。

4 总结

通过本节实验，你已经成功将一个 NLP 模型(BERT)以TorchScript格式部署到Triton Inference Server，并进一步探索了通过ONNX和TensorRT格式实现的部署优化。在下一篇文章中，我们将学习如何对模型本身进行优化，并以更高效的方式进行部署…