训练后保存的模型产物,例如 PyTorch 的 .pt 文件或 TensorFlow SavedModel 文件,主要表示模型的架构和学习到的权重。然而,它并非是为高性能推理而优化的可执行文件。为了让训练好的模型能够满足严格的延迟和吞吐量服务水平目标(SLOs),并成为可用于生产的服务,我们必须进行一系列优化步骤。这些步骤将模型的计算图转换为能在目标硬件上高效运行的格式。在这方面,两个最重要的工具是开放神经网络交换(ONNX)格式及其关联的运行时,以及 NVIDIA 的 TensorRT。它们共同提供了一个高效流程,用于将特定框架的模型转换为高性能推理引擎。ONNX 的解耦能力许多优化流程中的第一步是将模型从其原始训练框架导出为 ONNX 格式。ONNX 是一种开放标准格式,用于表示机器学习模型。可以将其视为一种中间表示(IR),它将模型的架构与其创建所在的框架解耦。这种可移植性是其主要益处。一旦模型转换为 ONNX 格式,它就可以由任何兼容的运行时或编译器运行,使您摆脱供应商锁定,并提供一致的部署目标。digraph G { rankdir=TB; node [shape=box, style="rounded,filled", fillcolor="#a5d8ff"]; edge [color="#495057"]; subgraph cluster_frameworks { label="训练框架"; style=dashed; color="#868e96"; bgcolor="#f8f9fa"; PyTorch [label="PyTorch\n(.pt)", fillcolor="#ffd8a8"]; TensorFlow [label="TensorFlow\n(SavedModel)", fillcolor="#ffc9c9"]; } subgraph cluster_runtimes { label="推理环境"; style=dashed; color="#868e96"; bgcolor="#f8f9fa"; ORT [label="ONNX Runtime\n(CPU, GPU)", fillcolor="#b2f2bb"]; TRT [label="NVIDIA TensorRT\n(NVIDIA GPU)", fillcolor="#96f2d7"]; Other [label="其他运行时\n(边缘设备, 移动设备)", fillcolor="#ced4da"]; } ONNX [label="ONNX 模型\n(.onnx)", shape=cylinder, fillcolor="#eebefa", height=1.2]; PyTorch -> ONNX; TensorFlow -> ONNX; ONNX -> ORT; ONNX -> TRT; ONNX -> Other; }ONNX 生态系统。模型从各种训练框架导出到通用的 ONNX 格式,然后可以被多个推理引擎和编译器使用。导出模型通常是一个简单的过程。例如,在 PyTorch 中,您可以使用 torch.onnx.export() 函数:import torch import torchvision # 加载预训练模型 model = torchvision.models.resnet50(pretrained=True) model.eval() # 创建具有正确尺寸的虚拟输入 dummy_input = torch.randn(1, 3, 224, 224) # 将模型导出为 ONNX 格式 torch.onnx.export(model, dummy_input, "resnet50.onnx", input_names=['input'], output_names=['output'], dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})dynamic_axes 参数对于推理服务器尤为重要,因为它允许模型接受不同大小的批次,这一特性对于动态批处理等技术来说必不可少。使用 ONNX Runtime 进行加速尽管 ONNX 定义了格式,ONNX Runtime 是一个为执行这些模型而构建的高性能推理引擎。它不仅仅是一个简单的解释器。加载 .onnx 文件后,ONNX Runtime 会应用一系列与硬件无关的图优化,包括:常量折叠: 预先计算图中仅依赖于常量输入的部分。节点消除: 移除冗余节点,例如恒等操作或 Dropout(推理时不需要)。节点融合: 将多个简单操作合并为一个更复杂的操作。例如,将 MatMul 操作与随后的 Add(用于偏置)融合为一个 FusedMatMul。ONNX Runtime 的一个重要特性是它使用执行提供程序(EPs)。EP 是一种后端,允许 ONNX Runtime 将图执行委托给专门的硬件库。您可以使用默认的 EP 在 CPU 上运行相同的 ONNX 模型,或者通过指定 CUDA EP 或 TensorRT EP 在 NVIDIA GPU 上对其进行加速,所有这些都无需更改模型文件。这种架构在性能和跨平台兼容性之间取得了很好的平衡。使用 NVIDIA TensorRT 实现最高性能对于部署在 NVIDIA GPU 上的工作负载,TensorRT 作为深度学习编译器和运行时,能提供最高水平的性能。尽管带有 CUDA EP 的 ONNX Runtime 在 GPU 上执行操作的通用版本,但 TensorRT 走得更远。它接受模型(通常是 ONNX 格式),并执行深度硬件特定优化,以生成序列化的“引擎”文件。这个过程计算成本高且需预先完成,但生成的引擎是为特定的 GPU 架构(例如 Ampere A100、Hopper H100)和特定的精度量身定制的。TensorRT 的主要优化包括:1. 图与层融合TensorRT 积极地融合层,以最大程度地减少内存带宽使用和内核启动开销。它可以将卷积、偏置加法和 ReLU 激活等连续操作组合成一个“CBR”内核。这意味着这些层之间的中间数据永远不需要写入和读取全局 GPU 内存,从而显著降低了延迟。digraph G { graph [rankdir=LR]; node [shape=box, style=rounded, fontname="sans-serif"]; edge [fontname="sans-serif"]; subgraph cluster_before { label="TensorRT 融合前"; style=filled; color="#f8f9fa"; node [style=filled, fillcolor="#ffc9c9"]; A [label="输入张量"]; B [label="卷积"]; C [label="偏置加法"]; D [label="ReLU"]; E [label="输出张量"]; A -> B [label="读取"]; B -> C [label="写入/读取"]; C -> D [label="写入/读取"]; D -> E [label="写入"]; } subgraph cluster_after { label="TensorRT 融合后"; style=filled; color="#f8f9fa"; node [style=filled, fillcolor="#b2f2bb"]; A_f [label="输入张量"]; Fused [label="融合内核\n(卷积 + 偏置 + ReLU)"]; E_f [label="输出张量"]; A_f -> Fused [label="读取"]; Fused -> E_f [label="写入"]; } }TensorRT 的层融合通过将多个节点组合成一个优化的内核,减少了内存操作。2. 内核自动调优NVIDIA GPU 包含数千个核心,并且通常有许多不同的算法(内核)来实现像卷积这样的单一操作。TensorRT 维护着这些内核的库,并在引擎构建过程中,它会对模型中的层测试多种实现方案。然后它会为特定的张量维度和目标 GPU 选择最快的内核,从而为您的模型高效创建定制的计算路径。3. 精度校准TensorRT 是执行训练后量化(PTQ)的主要工具。它可以将 32 位浮点(FP32)模型转换为使用 8 位整数(INT8)。这不仅将模型的内存占用减少了 4 倍,还在现代 NVIDIA GPU 上使用了专门的 Tensor Cores,以显著提升性能。为了在不大幅降低准确性的情况下实现这一点,TensorRT 使用校准过程,在该过程中,它会在一小部分代表性数据样本上运行 FP32 模型,以测量激活值的分布。然后它使用这些信息来确定将浮点范围转换为有限的 INT8 范围的最佳缩放因子。这种转换基于仿射映射: $$ v_{float} \approx S \cdot (v_{quant} - Z) $$ $v_{float}$ 是实际值,$v_{quant}$ 是量化后的整数值,$S$ 是缩放因子,而 $Z$ 是零点。TensorRT 的校准过程旨在找到最小化信息损失的最佳 $S$ 和 $Z$。选择您的优化策略ONNX Runtime 和 TensorRT 之间的选择取决于您的具体性能要求和部署限制。在以下情况使用 ONNX Runtime:您需要在广泛的硬件上实现可移植性,包括 CPU 和非 NVIDIA GPU。“足够好”的 GPU 性能可以接受,且开发简易性是优先考虑的。您想要一个可以在多个环境中部署的单一模型产物。在以下情况使用 TensorRT:您部署在 NVIDIA GPU 上,并且需要绝对最低的延迟和最高的吞吐量。您可以承担离线构建时间来生成硬件特定的引擎。管理针对不同 GPU 架构的不同引擎文件的操作开销可以接受。一种常见且有效的策略是将它们结合使用。标准的工作流程包括将模型导出为 ONNX 格式,然后将 ONNX 文件作为 TensorRT 构建过程的输入。这使得 ONNX 成为一个稳定且不依赖于框架的起点,以便进行 TensorRT 积极的、平台特定的优化。{"layout":{"title":{"text":"GPU 推理延迟比较"},"xaxis":{"title":{"text":"模型"}},"yaxis":{"title":{"text":"延迟 (毫秒)"},"type":"log"},"barmode":"group","font":{"family":"sans-serif"},"plot_bgcolor":"#f8f9fa","paper_bgcolor":"#FFFFFF"},"data":[{"x":["ResNet-50","BERT-Base"],"y":[12.5,28.1],"name":"PyTorch Eager (FP32)","type":"bar","marker":{"color":"#fa5252"}},{"x":["ResNet-50","BERT-Base"],"y":[7.2,15.4],"name":"ONNX Runtime (FP32)","type":"bar","marker":{"color":"#339af0"}},{"x":["ResNet-50","BERT-Base"],"y":[2.1,5.8],"name":"TensorRT (FP32)","type":"bar","marker":{"color":"#51cf66"}},{"x":["ResNet-50","BERT-Base"],"y":[0.9,2.3],"name":"TensorRT (INT8)","type":"bar","marker":{"color":"#fcc419"}}]}常见模型在不同执行后端下的示例性延迟比较。注意对数刻度。优化带来的性能提升显著,其中 INT8 精度的 TensorRT 提供了最低的延迟。