📚 课程目录第4章模型转换流程

4. 模型转换流程：从PyTorch到ONNX再到TensorRT引擎的完整转换步骤

4.1 核心原理：为什么需要两段式转换

在生产环境中，我们很少直接将PyTorch模型交给TensorRT部署。PyTorch模型本质上是Python运行时中的动态计算图，其算子粒度、内存布局和内核调度方式与TensorRT的优化目标存在显著差异。直接转换往往导致大量算子不被支持，或被迫回退到低效的插件实现。因此，业界标准路径是：PyTorch → ONNX → TensorRT引擎。ONNX在此过程中充当“中间表示层”，它屏蔽了PyTorch的动态控制流和Python对象，将模型固化为静态计算图，同时保留了完整的算子语义和维度信息，为TensorRT的层融合、精度校准和内核自动调优提供干净的输入。

关键认知：ONNX不是“最终产物”，而是“翻译中间人”。它的核心价值在于：① 将动态图固化为静态图；② 提供标准化的算子集；③ 允许我们在导出阶段就发现不支持的算子，而不是等到TensorRT构建引擎时才报错。

4.2 方法流程：五步完成模型转换

完整的转换流程可以拆解为五个步骤，每一步都有明确的输入、输出和验证手段。下面以生产中最常见的动态尺寸BERT模型为例，给出可执行的转换路径。

步骤一：PyTorch模型导出为ONNX

导出前必须将模型切换到eval()模式，并确保输入张量使用正确的设备。对于动态维度，需要在torch.onnx.export中通过dynamic_axes参数声明。以下是一个经过生产验证的导出配置：

import torch
from transformers import BertModel

model = BertModel.from_pretrained("bert-base-uncased").eval().cuda()
dummy_input = torch.randint(0, 30000, (1, 128), dtype=torch.long).cuda()
dummy_mask = torch.ones(1, 128, dtype=torch.long).cuda()

torch.onnx.export(
    model,
    (dummy_input, dummy_mask),
    "bert_base.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["last_hidden_state", "pooler_output"],
    dynamic_axes={
        "input_ids": {0: "batch_size", 1: "seq_len"},
        "attention_mask": {0: "batch_size", 1: "seq_len"},
        "last_hidden_state": {0: "batch_size", 1: "seq_len"},
        "pooler_output": {0: "batch_size"}
    },
    opset_version=17,
    do_constant_folding=True
)
print("ONNX导出完成")

常见陷阱：导出时忘记设置do_constant_folding=True会导致ONNX图中残留大量冗余常量节点，增加后续TensorRT解析负担。另外，opset_version并非越高越好，需确认TensorRT版本支持的ONNX算子集上限（TensorRT 8.6支持opset 17，TensorRT 10.x支持opset 19+）。

步骤二：ONNX图验证与算子检查

导出后立即用onnxruntime进行推理验证，确保ONNX模型输出与PyTorch原始输出一致（误差小于1e-4）。同时使用onnx.checker检查图结构合法性：

import onnx
import onnxruntime as ort
import numpy as np

# 检查ONNX图结构
onnx_model = onnx.load("bert_base.onnx")
onnx.checker.check_model(onnx_model)

# 使用onnxruntime验证输出
sess = ort.InferenceSession("bert_base.onnx", providers=["CUDAExecutionProvider"])
test_input = np.random.randint(0, 30000, (2, 64), dtype=np.int64)
test_mask = np.ones((2, 64), dtype=np.int64)
ort_out = sess.run(None, {"input_ids": test_input, "attention_mask": test_mask})

# 与PyTorch输出对比
with torch.no_grad():
    torch_out = model(torch.from_numpy(test_input).cuda(),
                      torch.from_numpy(test_mask).cuda())
    torch_out = [t.cpu().numpy() for t in torch_out]

for i, (a, b) in enumerate(zip(ort_out, torch_out)):
    diff = np.max(np.abs(a - b))
    print(f"输出{i}最大差异: {diff:.2e}")
    assert diff < 1e-4, "ONNX输出与PyTorch不一致！"

步骤三：TensorRT引擎构建

推荐使用trtexec命令行工具进行初步验证，它能快速暴露不支持的算子或维度问题。对于动态形状模型，必须显式指定--minShapes、--optShapes和--maxShapes三个档位：

trtexec \
  --onnx=bert_base.onnx \
  --saveEngine=bert_base_fp16.engine \
  --fp16 \
  --minShapes=input_ids:1x64,attention_mask:1x64 \
  --optShapes=input_ids:8x128,attention_mask:8x128 \
  --maxShapes=input_ids:32x256,attention_mask:32x256 \
  --workspace=4096 \
  --verbose

构建过程中重点关注日志中的Layer fusion信息和Total Host Memory统计。如果出现UNSUPPORTED_OPERATOR错误，需要回到ONNX层面进行算子替换或拆解。

步骤四：引擎精度与性能验证

引擎构建完成后，必须使用与训练时相同的数据分布进行精度验证。建议准备一个包含100~500条样本的验证集，分别计算PyTorch、ONNX Runtime和TensorRT引擎的输出，统计余弦相似度和最大绝对误差：

对比项	PyTorch (FP32)	ONNX Runtime	TensorRT FP16	TensorRT INT8
平均延迟 (ms)	12.5	10.2	3.8	2.1
P99延迟 (ms)	18.3	15.7	6.2	4.0
余弦相似度	1.0	0.99998	0.99991	0.99872
最大绝对误差	0	3.2e-5	2.1e-3	7.8e-2

如果FP16精度不满足业务要求（例如NLP任务中误差超过1e-2），可考虑对特定层使用--layerPrecisions进行混合精度控制，或回退到FP32构建。

步骤五：引擎序列化与部署集成

验证通过的引擎文件（.engine）可以直接序列化保存，部署时通过TensorRT的deserialize_cuda_engine加载。注意引擎文件与GPU架构绑定，不同型号的GPU（如A10与A100）不能共用同一引擎文件，必须在目标机型上重新构建。

import tensorrt as trt

# 反序列化引擎
with open("bert_base_fp16.engine", "rb") as f:
    engine_bytes = f.read()

runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
engine = runtime.deserialize_cuda_engine(engine_bytes)

# 创建执行上下文
context = engine.create_execution_context()
# 设置动态形状的实际尺寸
context.set_input_shape("input_ids", (8, 128))
context.set_input_shape("attention_mask", (8, 128))

4.3 案例：BERT模型转换中的动态维度处理

以生产环境中的在线问答系统为例，请求的文本长度从几十到几百不等。如果固定输入维度为256，短文本请求会浪费大量算力；如果固定为64，长文本请求则无法处理。因此必须使用动态形状（Dynamic Shape）能力。

在上述步骤一的导出代码中，我们通过dynamic_axes将batch_size和seq_len均声明为动态。在步骤三的trtexec命令中，三个档位的设置决定了引擎能接受的形状范围：

minShapes：1×64，对应最小请求（单条短文本）
optShapes：8×128，对应最常见请求（8条并发，平均长度128）
maxShapes：32×256，对应最大请求（32条并发，最长文本）

实际部署时，每次推理前必须调用context.set_input_shape()设置当前批次的实际形状。如果输入形状超出maxShapes范围，TensorRT会抛出Assertion failed错误。因此，生产环境必须在API网关层做文本长度截断或分批策略，确保进入引擎的请求不超过预设上限。

  性能对比数据：在A10 GPU上，固定形状（8×128）的BERT引擎延迟为3.2ms；动态形状引擎在同样形状下延迟为3.8ms，性能损失约18%。但动态形状带来的灵活性使得GPU利用率提升了约40%，因为短文本请求不再被强制填充到256长度。对于延迟敏感型业务，建议同时构建固定形状引擎和动态形状引擎，通过路由策略按需分发。

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

根据多个生产项目的故障复盘，以下问题占据了模型转换失败案例的80%以上：

故障现象	根因分析	解决方案
ONNX导出时报`Unsupported operator: aten::nonzero`	PyTorch模型中含有数据相关的控制流（如`torch.nonzero`、`torch.where`条件分支）	将动态控制流改为固定逻辑，或使用`torch.onnx.is_onnx_backend`检查；必要时拆分为多个ONNX子图
TensorRT构建时报`UNSUPPORTED_OPERATOR: GatherElements`	ONNX算子版本超出TensorRT支持范围，或该算子未被TensorRT实现	升级TensorRT版本；或修改模型结构，用`Gather`+`Unsqueeze`组合替代
引擎构建成功但推理结果全为NaN	FP16模式下某些层出现数值溢出，常见于LayerNorm或Softmax的中间结果	对这些层强制使用FP32：`--layerPrecisions=layer_name:fp32`；或使用`--fp16`时增加`--strictTypes`检查
动态形状引擎在运行时崩溃	实际输入形状超出构建时设置的`maxShapes`范围	在推理服务入口增加形状校验；或重新构建引擎并扩大`maxShapes`

最重要的复盘结论：模型转换不是“一键完成”的操作，而是一个需要反复迭代的工程过程。建议在项目初期就搭建自动化的转换验证流水线，将PyTorch输出、ONNX输出和TensorRT输出三方对比作为CI/CD的必选检查项。任何模型结构变更后，都必须重新跑完整转换流程，不能只更新权重文件。

最后需要强调的是，转换流程中的每一步都应有明确的产物和验证标准：PyTorch导出后检查ONNX图结构，ONNX验证后检查数值一致性，TensorRT构建后检查精度和性能。只有每一层都严格把关，才能避免将问题带到生产环境，减少线上故障排查的代价。