4. 模型转换与优化：将PyTorch模型转换为TensorRT引擎、精度校准（FP16/INT8）、动态形状支持

4.1 核心原理：从PyTorch图到TensorRT引擎的转换链路

PyTorch模型默认以动态图（Eager Mode）方式执行，每个算子在前向传播时被逐个调度，框架层开销大且难以做跨算子融合。TensorRT则要求先获得静态计算图，再通过层融合、内核自动调优、显存复用等手段生成高度优化的引擎文件。整个转换链路可概括为：.pt/.pth → ONNX → TensorRT Engine（.engine/.plan）。

转换的本质是“语义等价、结构重写”：ONNX作为中间表示（IR）保留了算子的数学语义，TensorRT解析ONNX后，将多个算子融合为单一内核（如将Conv+Bias+ReLU融合为CBR），并针对目标GPU架构（如SM版本、显存带宽）生成最优的CUDA内核与执行计划。值得注意的是，TensorRT并非逐算子翻译，而是构建一个执行图（Execution Graph），其中节点是融合后的内核，边是显存中的张量。

PyTorch 模型 动态图 / Eager 权重 .pt/.pth 含控制流、Python 循环等动态结构 torch.onnx.export 导出 ONNX 图 静态计算图 算子级 IR 动态维度用 symbolic dim onnxruntime 验证 解析 TensorRT 引擎 .engine / .plan 融合内核 + CUDA 优化 显存池 / 自动调优 支持 FP16/INT8 动态形状 Profile 关键：ONNX 是语义等价中间层，TensorRT 是结构重写与内核优化层

上图展示了转换链路的三个关键阶段。实际项目中，ONNX导出往往是最容易出问题的环节——PyTorch的torch.onnx.export依赖符号追踪（Tracing），它只能记录实际执行过的张量操作，遇到if分支、for循环依赖张量值时，会静态展开或报错。因此，转换前必须将模型推理路径固定为数据无关的控制流。

4.2 方法流程：从PyTorch到TensorRT的完整转换步骤

下面给出一个经过生产验证的转换流程，分为五个步骤。每一步都有明确的输入、输出和验证手段。

  1. 模型冻结与预处理：将PyTorch模型设为eval()模式，并确保所有BatchNorm层被折叠（可通过torch.jit.script或直接导出ONNX时自动折叠）。移除Dropout等训练专用层。
  2. 导出ONNX：使用torch.onnx.export，设置opset_version=17（或更高），并指定dynamic_axes以保留动态维度。
  3. ONNX验证与简化：用onnxruntime加载ONNX并跑通推理，对比PyTorch输出（误差<1e-4）。使用onnx-simplifier消除冗余节点（如Identity、Shape-Gather序列）。
  4. 构建TensorRT引擎：通过trtexec或Python API构建引擎。关键参数包括--fp16、--int8、--minShapes/--optShapes/--maxShapes。
  5. 引擎验证与性能测试：加载引擎，用真实输入张量跑推理，对比输出与ONNX/PyTorch的误差，并测量延迟与吞吐。

以下是一个典型的trtexec命令行示例，展示了动态形状与INT8校准的配置方式：

trtexec \
  --onnx=model.onnx \
  --saveEngine=model_fp16.engine \
  --fp16 \
  --minShapes=input:1x3x224x224 \
  --optShapes=input:8x3x224x224 \
  --maxShapes=input:32x3x224x224 \
  --calib=calibration.cache \
  --int8 \
  --verbose

其中--calib指向校准缓存文件。若未提供缓存，TensorRT会使用随机数据或用户提供的校准数据生成。生产环境必须使用代表性数据集进行校准，否则INT8精度可能严重劣化。

4.3 精度校准：FP16与INT8的量化策略

FP16转换相对简单，只需将权重和激活从FP32降为FP16，TensorRT自动处理精度溢出风险。但INT8需要校准（Calibration），其核心是统计每一层激活值的分布，从而确定缩放因子（Scale），将浮点张量映射到[-128, 127]的整数范围。

TensorRT提供三种校准算法：MinMax、Entropy（默认）和Legacy。其中Entropy（基于KL散度）通过最小化量化前后信息损失来选择阈值，对大多数视觉模型效果最好。校准过程需要前向传播校准数据集（通常500~2000张代表性图片），记录每层激活的直方图，然后计算最优截断阈值。

精度模式 加速比（相对FP32） 精度损失 适用场景 关键注意事项
FP32 1.0x 无 精度敏感、小模型 显存占用高
FP16 1.5~2.5x 极小（通常<0.5%） 大多数推理场景 需检查是否有溢出（如大数值梯度）
INT8 3~5x 1%~3%（视校准质量） 高吞吐、边缘部署 校准集必须代表真实数据分布

校准过程中常见的误区是使用训练集而非验证集进行校准。训练集可能包含标签噪声或数据增强后的分布偏移，导致校准统计失真。推荐从验证集中随机抽取子集，并确保覆盖各类别、各种光照/角度条件。

4.4 动态形状支持：Profile与显存优化

动态形状（Dynamic Shape）允许引擎在运行时接受不同batch size或不同分辨率的输入。TensorRT通过Profile机制实现：每个Profile定义一组min/opt/max形状，引擎在opt形状处做最优内核选择，在min和max之间动态调整显存分配。

使用Python API构建动态形状引擎时，需要显式创建Profile并设置形状范围：

import tensorrt as trt

builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
network = builder.create_network(1 << trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
parser = trt.OnnxParser(network, trt.Logger(trt.Logger.WARNING))

with open("model.onnx", "rb") as f:
    parser.parse(f.read())

profile = builder.create_optimization_profile()
input_tensor = network.get_input(0)
profile.set_shape(input_tensor.name, 
                  min=(1, 3, 224, 224),
                  opt=(8, 3, 224, 224),
                  max=(32, 3, 224, 224))
config = builder.create_builder_config()
config.add_optimization_profile(profile)
config.set_flag(trt.BuilderFlag.FP16)

engine = builder.build_serialized_network(network, config)

动态形状带来的显存挑战在于：TensorRT必须为max形状预留显存池，即使实际推理时batch size很小。这会导致显存利用率下降。优化策略包括：

  • 将max形状设置为业务上限，而非硬件上限。
  • 使用set_memory_pool_limit控制显存池上限，超出时TensorRT会回退到更保守的分配策略。
  • 若业务中batch size变化不频繁，可构建多个静态引擎（如batch=1、4、8），在运行时切换，避免动态形状的显存浪费。
实践建议：对于在线推理服务，推荐使用固定batch size + 动态分辨率的组合。例如，batch固定为1，但输入分辨率允许在[224, 512]范围内变化。这样既保留了动态形状的灵活性，又减少了显存池的浪费，同时内核选择更稳定。

4.5 案例：将ResNet-50转换为INT8引擎并验证精度

假设我们有一个在ImageNet上预训练的ResNet-50 PyTorch模型，目标是在T4 GPU上实现INT8推理，batch size动态范围1~16。以下是完整流程：

第一步：导出ONNX并验证

import torch
import torchvision.models as models

model = models.resnet50(pretrained=True).eval().cuda()
dummy_input = torch.randn(1, 3, 224, 224).cuda()

torch.onnx.export(
    model, dummy_input, "resnet50.onnx",
    input_names=["input"], output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
    opset_version=17
)

第二步：准备校准数据集。从ImageNet验证集中随机抽取1000张图片，预处理为与训练一致的格式（Resize到256、CenterCrop到224、归一化）。

第三步：使用trtexec构建INT8引擎。校准缓存生成后，后续构建可直接复用calibration.cache，无需重新校准。

第四步：精度对比。用100张验证图片分别跑PyTorch FP32和TensorRT INT8，计算Top-1准确率差异。通常INT8的Top-1下降应控制在1%以内。若超过2%，需检查校准集质量或考虑对敏感层跳过量化（通过--layerPrecisions指定）。

常见陷阱：当模型包含动态数值范围大的层（如检测头的回归分支、Embedding层）时，全局INT8量化可能导致精度骤降。此时应使用逐层精度控制：为这些层单独指定FP16精度，其余层保持INT8。TensorRT 8.6+支持set_precision接口，可在网络构建时逐层设置。

4.6 误区与复盘：转换失败的常见根因

根据多个生产项目的经验，模型转换失败或精度劣化通常源于以下五个方面：

现象 根因 解决方案
ONNX导出报错 模型含torch.where、torch.topk等动态控制流 用torch.jit.script或重写为固定逻辑
引擎构建成功但输出NaN FP16下数值溢出（如大梯度、大激活值） 检查是否有exp/pow操作，改用FP32或添加精度保护层
INT8精度下降超过5% 校准集与真实数据分布不匹配 重新采样校准集，增加数据多样性
动态形状推理报错 输入尺寸超出Profile的min/max范围 在服务端做尺寸裁剪或填充，确保落在范围内
引擎加载慢（>5s） 引擎文件过大或反序列化时重新分配显存 使用cudaGraph捕获或提前预热显存池

复盘时建议建立转换验证清单：每次转换后，至少对比3个不同输入尺寸下的输出张量（FP32 vs FP16 vs INT8），计算最大绝对误差和余弦相似度。只有当余弦相似度>0.999且最大绝对误差在可接受范围时，才可部署到生产环境。

最后需要强调的是，TensorRT引擎是硬件绑定的——在T4上构建的引擎无法在A100上运行。因此，CI/CD流水线中必须为每种GPU型号单独构建引擎，并缓存到对象存储中，避免每次部署时重复构建。

← 上一章 📖 返回目录 下一章 →