在本系列文章的这一部分,我们将从模型训练与优化,转向模型在实际应用中的落地:学习如何导出训练好的模型,使用 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 的核心部分,负责调度、执行推理,并返回结果。
- Model Management(模型管理)
Triton 会从模型仓库加载模型,并管理其生命周期。你可以指定多个模型同时部署,支持版本控制。 - Inference Request / Response(推理请求与响应)
请求进入服务器后被解析,分发给对应模型的执行队列,然后进行推理处理,最后返回响应。 - Per-model Scheduler Queues(模型调度队列)
每个模型有自己的调度队列,用于支持动态批处理(Dynamic Batching)、请求排序等调度策略,提高吞吐量。 - Framework Backends(框架后端)
Triton 支持多种深度学习框架的模型推理,图中列出了主要支持的格式和框架:TensorRT
、TensorFlow
、ONNX
、PyTorch
和自定义后端(Custom)
- Scheduler(调度器)
调度器根据资源利用率、请求优先级等因素分发请求到后端推理引擎。 - 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引擎等方式,以不同的格式和后端执行方式运行这些模型。
-
PyTorch JIT / TorchScript
TorchScript 是 PyTorch 官方支持的模型导出格式,可将动态图转换为可序列化、可部署的静态图。它适合直接将 PyTorch 模型导出为部署版本,无需转换框架。Triton 可以通过
libtorch
后端加载和执行 TorchScript 模型,是最接近 PyTorch 原生方式的部署方案,部署过程简单,调试也较为直观。 -
ONNX Runtime
ONNX(
Open Neural Network Exchange
)是一个开放的模型交换格式,适用于在不同框架间共享模型。PyTorch 模型可以被导出为 ONNX 格式,并通过 ONNX Runtime 高效推理。此方式适合需要跨平台、跨框架兼容性的部署场景,同时支持在 CPU 或 GPU 上运行,易于移植。 -
ONNX-TensorRT
在导出为 ONNX 格式后,可以进一步利用 TensorRT 对模型进行图优化、算子融合和精度压缩(如 FP16、INT8),生成高效的TensorRT引擎。这种方式结合了ONNX的兼容性和TensorRT的高性能,适用于需要高吞吐、低延迟的 GPU 推理场景,是生产部署中常见的优化路径。
-
TensorRT
TensorRT是NVIDIA推出的专用推理引擎,支持直接加载经过转换的 TensorRT 引擎文件(如从 PyTorch 通过中间工具转化)。它提供了最极致的性能优化,包括动态批处理、精度量化等,但转换过程相对复杂,适用于对性能要求极高、可控性强的部署环境。
虽然我们当前用的是 PyTorch + BERT 模型,但Triton还能支持很多别的格式和框架,因此在一个统一的平台下可以部署多个来源的模型。
- TensorFlow GraphDef
TensorFlow 的早期模型格式,是用.pb
文件保存的静态图(graph definition
)。如果你使用 TensorFlow 1.x 训练模型,可以直接导出为 GraphDef,然后在 Triton 中部署。 - TensorFlow SavedModel
TensorFlow 2.x推荐使用的模型格式,支持保存计算图和权重。SavedModel更完整、可扩展,适合现代 TensorFlow 项目,也是 Triton 原生支持的一种格式。 - Caffe2 导出格式
虽然Caffe2现在较少使用(已被PyTorch整合),但Triton仍然支持其模型部署。适用于那些历史上使用 Caffe2 构建的工业项目。 - 自定义模型(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 注册一个新模型需执行以下步骤:
- 在模型仓库中创建一个新的模型文件夹。该文件夹名称应与希望对外暴露的服务名称一致。
- 在模型文件夹中创建
config.pbtxt
文件,提供基本的服务配置。 - 在模型文件夹中创建至少一个模型版本的子文件夹。子文件夹名称为模型版本号,您可以托管多个版本的同一模型。
接下来,我们将演示如何将模型导出到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
格式实现的部署优化。在下一篇文章中,我们将学习如何对模型本身进行优化,并以更高效的方式进行部署…