4. 模型转换与ONNX导出实战

在 Jetson 上部署模型，TensorRT 并不直接读取 PyTorch 的 .pt 或 TensorFlow 的 SavedModel，而是以 ONNX 作为中间表示。ONNX 是一张静态计算图：节点是算子，边是张量，每个张量都带有形状和数据类型。导出质量直接决定后续 trtexec 与 TensorRT 能否顺利解析、能否启用 FP16/INT8、以及推理时是否需要频繁重建引擎。本节围绕“导出—优化—验证”这条主线，讲清 PyTorch/TensorFlow 的导出方法、动态轴处理、常见报错与修复，以及用 trtexec 快速验证的完整过程。

一、PyTorch 导出 ONNX

PyTorch 使用 torch.onnx.export 做追踪式导出。核心是提供一份与真实输入形状、类型一致的 dummy input，模型会按该输入执行一次前向，记录算子轨迹。以图像分类为例：

import torch, torchvision

model = torchvision.models.resnet18(weights="IMAGENET1K_V1").eval().cuda()
dummy = torch.randn(1, 3, 224, 224, device="cuda")

torch.onnx.export(
    model, dummy, "resnet18.onnx",
    input_names=["input"], output_names=["logits"],
    opset_version=17,
    dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
    do_constant_folding=True,
)

几个关键点：model.eval() 必须调用，否则 BatchNorm、Dropout 会以训练模式导出，产生多余的 BatchNormalization 训练分支；opset_version 建议 13～17，TensorRT 8.x 对 17 支持良好，过低会缺少 LayerNormalization 等算子，过高可能不被旧版 TensorRT 识别；do_constant_folding=True 会把可静态计算的子图折叠成常量，减小图体积。

二、TensorFlow 导出 ONNX

TensorFlow 侧推荐使用 tf2onnx。对 Keras 模型可直接从内存转换：

import tf2onnx, tensorflow as tf

model = tf.keras.applications.MobileNetV2(weights="imagenet")
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
model.output_names = ["logits"]

tf2onnx.convert.from_keras(
    model, input_signature=spec, opset=17, output_path="mobilenetv2.onnx"
)

若已有 SavedModel，则用 python -m tf2onnx.convert --saved-model ./saved_model --output model.onnx --opset 17。注意 TensorFlow 默认是 NHWC 布局，而 TensorRT 在 GPU 上更偏好 NCHW，转换时可通过 --inputs input:0[1,224,224,3] 显式指定形状，或在图中插入 Transpose。若模型含 tf.nn.top_k、动态 tf.shape 等控制流，tf2onnx 可能生成 Loop/If 子图，TensorRT 对这类结构的支持有限，应尽量在导出前用静态形状替换。

三、动态轴与动态形状处理

动态轴是导出中最容易出问题的一环。它让 batch、序列长度、图像尺寸在推理时可变化，但代价是 TensorRT 需要为每个 shape 建立优化配置（Optimization Profile）。导出时用 dynamic_axes 声明哪些维度可变，例如 NLP 模型：

dynamic_axes = {
    "input_ids": {0: "batch", 1: "seq_len"},
    "attention_mask": {0: "batch", 1: "seq_len"},
    "logits": {0: "batch", 1: "seq_len"},
}

导出后务必用 onnx.checker.check_model 校验，并用 onnx.shape_inference.infer_shapes 补全中间张量形状。若某维度在图中被 Reshape、Gather 等算子以常量方式引用，动态化会失败，此时需要改写模型，把硬编码的 view(-1, 512) 换成基于 shape 的动态计算。TensorRT 构建引擎时通过 --minShapes、--optShapes、--maxShapes 给出范围，optShapes 应贴近实际业务最常见的形状，否则性能会明显下降。

ONNX 导出与 trtexec 验证链路 源模型 PyTorch .pt TF SavedModel Keras Model eval() / 静态形状 导出工具 torch.onnx.export tf2onnx.convert opset 13~17 dynamic_axes 声明 ONNX 图 checker 校验 shape_inference 常量折叠 onnxsim 简化 trtexec 验证 --onnx=model.onnx --fp16 / --int8 --min/opt/maxShapes 输出 engine + 性能 动态轴声明示例 input_ids : {0:"batch", 1:"seq_len"} logits : {0:"batch", 1:"seq_len"} ✗ 图中硬编码 view(-1,512) 会破坏动态性 ✓ 改用 shape 动态计算后再 Reshape TensorRT Optimization Profile minShapes : input:1x3x224x224 optShapes : input:8x3x224x224 ← 贴近业务 maxShapes : input:32x3x224x224 opt 偏离实际形状 → 性能显著下降 常见导出错误与修复 ① 训练模式导出 → 调用 model.eval() ② 不支持的算子 → 升级 opset 或替换算子 ③ 动态轴冲突 → 检查 Reshape/Gather 常量维度 ④ NHWC/NCHW 不匹配 → 插入 Transpose ⑤ 控制流 Loop/If → 静态化或改写模型 ⑥ 精度下降 → 检查 FP16 溢出层

四、ONNX 图优化

导出后的图往往包含冗余节点，可用 onnxsim 做常量折叠、算子融合与死代码消除：

python -m onnxsim resnet18.onnx resnet18_sim.onnx

典型收益是消除导出时残留的 Identity、Cast、Transpose 链，把 Conv+BN+Relu 合并为更紧凑的结构。注意 onnxsim 会尝试固定动态维度，若模型需要动态 batch，应加 --dynamic-input-shape 参数。此外，TensorRT 自身在解析 ONNX 时也会做层融合，因此 ONNX 层面的优化主要是减小图体积、提升解析成功率，而非直接决定最终性能。

五、用 trtexec 快速验证

trtexec 是 TensorRT 自带的命令行工具，能在不写任何 C++/Python 代码的情况下完成解析、构建、推理与性能测试。基本用法：

trtexec --onnx=resnet18.onnx \
        --saveEngine=resnet18_fp16.engine \
        --fp16 \
        --minShapes=input:1x3x224x224 \
        --optShapes=input:8x3x224x224 \
        --maxShapes=input:32x3x224x224 \
        --workspace=2048 \
        --verbose

运行后重点关注几类输出：解析阶段若出现 Unsupported operator，说明该算子需用插件或改写；构建阶段会打印各层的精度与耗时，若某层回退到 FP32，会标注 fp32；性能阶段输出 Throughput（吞吐，单位 qps）与 GPU Compute Time（单次推理耗时）。判断是否达标，可对比 FP32 与 FP16 的吞吐提升，一般视觉模型 FP16 能带来 1.5～2 倍加速；若提升很小，需检查是否有大量层未融合或存在内存瓶颈。若只想验证 ONNX 能否被解析，可加 --skipInference 快速检查。

导出前先确认模型处于推理模式、输入形状与真实推理一致；导出后用 onnx.checker 与 onnxsim 各过一遍，再交给 trtexec，能避免绝大多数“解析失败”类问题。
动态形状并非越多越好。每个可变维度都会增加 TensorRT 的优化配置数量，导致构建时间变长、显存占用上升。若业务中 batch 固定，建议直接导出静态形状，性能最优。

整体流程可以概括为：源模型 → 导出 ONNX（声明动态轴）→ 校验与简化 → trtexec 构建引擎并测性能 → 根据日志定位问题并回到导出环节修复。掌握这条链路后，后续在 Python/C++ 中调用 TensorRT Runtime 时，只需复用已验证的 engine 与 profile 配置即可。

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