深度学习模型部署(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 支持多种深度学习框架的模型推理,图中列出了主要支持的格式和框架:TensorRTTensorFlowONNXPyTorch自定义后端(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运行模型
  • 使用TensorRTONNX转换为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 runtimeTensorRT的性能:

(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,并进一步探索了通过ONNXTensorRT格式实现的部署优化。在下一篇文章中,我们将学习如何对模型本身进行优化,并以更高效的方式进行部署…

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

tilblackout

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值