模型导出与格式转换：从PyTorch到TensorRT的实战之路

模型训练好只是第一步，真正要上线跑推理，格式转换这关必须过。我见过不少团队，模型在GPU上跑得飞快，一部署到嵌入式设备就卡成PPT——十有八九是中间格式没处理好。

今天咱们就聊聊PyTorch→ONNX→TensorRT这条链路。说白了，就是把一个"通用模型"变成"针对特定硬件优化过的模型"。你想想看，PyTorch模型里有很多动态操作，比如if分支、循环，这些在推理引擎里根本跑不了。所以我们需要一个中间格式来"固化"计算图。

3.1 PyTorch模型导出为ONNX

ONNX（Open Neural Network Exchange）是微软牵头搞的中间表示格式。它的核心价值在于：一次导出，多处部署。我最早接触ONNX是在2019年，当时踩了不少坑，现在把这些经验分享给你。

3.1.1 基础导出流程

PyTorch导出ONNX主要靠torch.onnx.export这个API。先看一个最简单的例子：

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(1, 32, 3)
        self.relu = nn.ReLU()
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(32, 10)

    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        x = self.pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = SimpleModel()
model.eval()

dummy_input = torch.randn(1, 1, 224, 224)

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

这里有几个关键点我要强调一下：

  • model.eval()：必须切到eval模式，否则BatchNorm、Dropout这些层会乱套
  • dummy_input：形状要和实际推理时一致，尤其是语音模型里的时间维度
  • dynamic_axes：语音任务里batch size和序列长度经常变化，这个参数必须配
  • opset_version：我建议用17或以上，支持的操作更多
注意：导出时如果遇到"Unsupported operator"错误，八成是模型里用了ONNX不支持的操作。比如torch.einsum、F.upsample_bilinear这些。解决办法是用torch.onnx.symbolic_helper注册自定义符号，或者改模型结构。

3.1.2 语音模型的特殊处理

语音模型和图像模型不太一样。我做过一个语音识别项目，模型里有torch.stft和torch.istft，这两个操作ONNX原生不支持。当时我折腾了两天，最后用torch.fx把计算图拆开，手动实现了STFT的ONNX版本。

常见的语音模型导出问题：

操作 问题 解决方案
torch.stft ONNX不支持复数运算 用torch.view_as_real拆成实部和虚部
torch.nn.GRU/LSTM 双向RNN的隐藏状态拼接 确保batch_first=True，导出时固定seq_len
torch.nn.Transformer 注意力掩码维度不匹配 用static shape导出，或者用ONNX的Pad操作模拟
torch.cat/stack 动态维度拼接 用onnx::ConcatFromSequence替代

3.2 ONNX模型优化与调试

导出ONNX只是第一步，接下来要优化。我习惯用onnx-simplifier和onnxruntime来做这件事。

3.2.1 模型简化

PyTorch导出的ONNX经常有一些冗余操作，比如多余的Reshape、Cast、Identity节点。这些节点不影响精度，但会增加推理延迟。

import onnx
from onnxsim import simplify

model = onnx.load("simple_model.onnx")
model_simp, check = simplify(model)

if check:
    onnx.save(model_simp, "simple_model_simplified.onnx")
    print("简化成功！")
else:
    print("简化失败，请检查模型结构")

简化后模型大小能减少20%-40%，推理速度提升10%-30%。我在一个语音唤醒模型上试过，从12MB降到了7MB，效果很明显。

3.2.2 精度验证

模型简化后一定要验证精度。我见过有人简化后模型输出全变成了NaN，还傻乎乎地拿去部署。

import onnxruntime as ort
import numpy as np

# 创建ONNX Runtime会话
ort_session = ort.InferenceSession("simple_model_simplified.onnx")

# 准备输入数据
input_data = np.random.randn(1, 1, 224, 224).astype(np.float32)

# 运行推理
outputs = ort_session.run(None, {"input": input_data})

# 对比PyTorch输出
with torch.no_grad():
    torch_outputs = model(torch.from_numpy(input_data)).numpy()

# 计算误差
diff = np.abs(outputs[0] - torch_outputs).max()
print(f"最大误差: {diff:.6f}")

if diff < 1e-4:
    print("精度验证通过！")
else:
    print("精度异常，请检查模型")
小技巧：我一般会对比多个输入样本的误差，而不是只测一个。因为有些操作（比如Softmax）在特定输入下误差会放大。建议至少测100个随机样本，看误差分布。

3.2.3 常见调试技巧

ONNX调试最头疼的是"模型能导出但推理结果不对"。我总结了一套排查流程：

  1. 检查输入输出名称：用onnx.load加载模型，打印graph的input和output
  2. 逐层对比：用onnxruntime的run方法指定输出中间层，和PyTorch逐层对比
  3. 检查数据类型：ONNX默认是float32，如果PyTorch用了float16或double，导出时可能出问题
  4. 检查动态轴：如果dynamic_axes配错了，推理时输入形状变化会导致崩溃

我曾经遇到一个坑：模型里用了torch.where，ONNX把它展开成了NonZero + Gather，结果在TensorRT上跑的时候，NonZero的输出形状不确定，导致后续节点报错。最后我改成torch.masked_select才解决。

3.3 TensorRT模型转换

TensorRT是NVIDIA的推理优化引擎。它能把ONNX模型编译成针对特定GPU优化的引擎文件（.engine或.plan）。

3.3.1 使用trtexec命令行工具

trtexec是TensorRT自带的命令行工具，用法很简单：

# 基本用法
trtexec --onnx=simple_model_simplified.onnx \
        --saveEngine=simple_model.engine \
        --workspace=4096 \
        --fp16 \
        --verbose

# 指定输入形状（动态batch）
trtexec --onnx=simple_model_simplified.onnx \
        --saveEngine=simple_model_dynamic.engine \
        --minShapes=input:1x1x224x224 \
        --optShapes=input:8x1x224x224 \
        --maxShapes=input:32x1x224x224 \
        --fp16

参数说明：

参数 作用 建议值
--onnx 输入ONNX模型路径 简化后的ONNX文件
--saveEngine 输出引擎文件路径 以.engine结尾
--workspace 工作空间大小（MB） 4096或更大
--fp16 启用FP16推理 如果GPU支持，强烈建议开启
--int8 启用INT8推理 需要校准数据集，精度可能下降
--minShapes/optShapes/maxShapes 动态形状范围 根据实际业务场景设置
核心要点：trtexec转换时，--workspace设得越大越好。我一般设8GB以上，否则复杂模型会报"out of memory"。另外，--verbose参数能打印详细的转换日志，方便排查问题。

3.3.2 语音模型的TensorRT优化

语音模型在TensorRT上有个特殊问题：序列长度变化。比如语音识别模型，输入音频长度从1秒到10秒不等。如果每个长度都重新编译引擎，那效率太低了。

我的做法是：

  • 固定序列长度：把输入音频padding到固定长度（比如8秒），用--optShapes指定这个长度
  • 动态形状：设置合理的--minShapes和--maxShapes，让TensorRT自动处理
  • 批量推理：把多个短音频拼成一个batch，减少引擎切换开销

我记得有个项目，语音端点检测模型输入长度变化很大。一开始用动态形状，结果TensorRT编译了3个小时还没完。后来改成固定长度padding，编译时间降到5分钟，推理速度还快了20%。

3.3.3 引擎验证与性能测试

转换完成后，一定要验证引擎的正确性和性能：

# 验证精度
trtexec --loadEngine=simple_model.engine \
        --shapes=input:1x1x224x224 \
        --noDataTransfers \
        --useSpinWait

# 性能测试
trtexec --loadEngine=simple_model.engine \
        --shapes=input:8x1x224x224 \
        --duration=10 \
        --avgRuns=100

性能测试输出里，重点关注End-to-End Host Latency和GPU Compute Time。前者包括数据拷贝时间，后者是纯计算时间。如果GPU Compute Time远小于End-to-End Latency，说明数据拷贝是瓶颈，可以考虑用CUDA Stream优化。

注意：TensorRT引擎是硬件相关的。在A100上编译的引擎，不能在T4上跑。部署时一定要在目标机器上重新编译，或者用TensorRT的"可移植性"特性（但有限制）。

3.4 本章知识体系

下面这张图总结了模型导出与格式转换的完整流程：

模型导出与格式转换流程 PyTorch模型 torch.onnx.export 动态轴 / opset_version ONNX模型 onnx-simplifier 精度验证 / 调试 TensorRT引擎 trtexec FP16/INT8优化 动态形状处理 min/opt/maxShapes 精度验证 逐层对比 / 误差分析 性能测试 延迟 / 吞吐量 常见问题与解决方案 • 导出失败：检查模型是否包含ONNX不支持的操作（如torch.stft、torch.einsum） • 精度异常：对比PyTorch和ONNX Runtime的逐层输出，定位误差来源 • 性能瓶颈：数据拷贝 vs 计算时间，考虑使用CUDA Stream或异步推理

整个流程环环相扣。PyTorch导出时埋下的坑，到TensorRT阶段才会暴露。所以我建议你：每一步都做验证，不要等到最后才发现问题。

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