模型转换入门：从PyTorch到ONNX，从ONNX到TensorRT，踩过的坑

模型转换这件事，说白了就是让模型换个语言说话。PyTorch讲的是Python，TensorRT讲的是CUDA和C++。中间那个ONNX，就是个翻译官。听起来简单吧？我刚开始做的时候也这么想，结果第一个模型就翻车了。

今天咱们就聊聊这个转换链条上的那些坑。我会把我在项目中踩过的雷，一个一个指给你看。

第一步：PyTorch → ONNX

这一步看着简单，其实暗藏玄机。PyTorch模型导出ONNX，核心就一个函数：torch.onnx.export()。但参数怎么配，学问大了。

核心代码示例：

import torch
import torch.onnx

# 假设你有一个训练好的模型
model = YourModel()
model.eval()

# 创建一个 dummy input
dummy_input = torch.randn(1, 3, 224, 224)

# 导出 ONNX
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 输入张量
    "model.onnx",        # 输出文件名
    export_params=True,  # 导出参数
    opset_version=11,    # ONNX opset 版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

嗯，这里要注意几个点。

第一，opset版本。我建议用11或更高。为什么？因为低版本不支持很多算子。我在项目中遇到过，用opset=9导出，结果TensorRT死活不认。后来发现是某个算子太新了，ONNX老版本不支持。

第二，dynamic_axes。这个参数很关键。你想想看，如果你的模型要处理不同batch size的输入，就必须设置动态轴。不然导出的ONNX就固定死了batch size=1。

小技巧：导出后一定要用 onnx.checker.check_model() 验证一下。我每次导出完都会跑一遍，能发现很多低级错误。

第二步：ONNX → TensorRT

这一步才是真正的考验。TensorRT的转换工具有两个：trtexec 命令行工具，和 TensorRT Python API。我个人习惯用Python API，因为可以精细控制。

TensorRT 转换代码：

import tensorrt as trt

# 创建 builder
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)

# 创建 network
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)

# 创建 parser
parser = trt.OnnxParser(network, logger)

# 解析 ONNX 模型
with open("model.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for error in range(parser.num_errors):
            print(parser.get_error(error))
        raise RuntimeError("ONNX 解析失败")

# 创建 config
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

# 构建 engine
serialized_engine = builder.build_serialized_network(network, config)

# 保存 engine
with open("model.engine", "wb") as f:
    f.write(serialized_engine)

这段代码看着简单，但坑都在细节里。

常见坑一：算子不支持

这是最头疼的问题。PyTorch里好好的算子，转到ONNX就变了样，再到TensorRT直接报错。

我曾经遇到过一个模型，里面用了 torch.nn.functional.grid_sample。这个算子在ONNX里映射成了 GridSample 节点，但TensorRT 8.x之前不支持。怎么办？

  • 方案一：升级TensorRT版本。新版本会支持更多算子。
  • 方案二：用插件。自己写一个TensorRT plugin，实现这个算子。
  • 方案三：改模型。把不支持的算子替换成支持的。比如用双线性插值替代grid_sample。

警告：不要盲目相信ONNX的算子支持列表。我建议你先用 onnxsim 简化模型，再用 polygraphy 做一次兼容性检查。

常见坑二：动态形状问题

ONNX导出了动态batch，但TensorRT处理起来很麻烦。你想想看，TensorRT需要做内存规划和算子融合，动态形状意味着这些优化不能提前做。

我的建议是：能固定就固定。如果业务场景允许，固定batch size和输入尺寸，性能会好很多。如果必须动态，那就用 optimization_profile 来指定形状范围。

动态形状配置示例：

# 创建 optimization profile
profile = builder.create_optimization_profile()

# 设置输入形状范围
profile.set_shape("input", 
    min_shape=(1, 3, 224, 224),
    opt_shape=(4, 3, 224, 224),
    max_shape=(8, 3, 224, 224)
)

config.add_optimization_profile(profile)

这里有个坑：opt_shape 要选你实际部署中最常见的batch size。TensorRT会针对这个尺寸做最优优化。选错了，性能会打折扣。

常见坑三：精度问题

模型转换后，精度下降是常有的事。我遇到过最夸张的一次，FP16推理比FP32差了5个点。排查了半天，发现是某个层的激活值范围太大，FP16存不下。

解决方案：

  • 使用FP32回退：对精度敏感的层，强制使用FP32。
  • 使用INT8量化：如果FP16精度不够，试试INT8。但需要校准数据集。
  • 使用精度校准：TensorRT提供了 Int8Calibrator，可以自动调整量化参数。

个人经验：我一般会先跑一次FP32推理，记录每层的激活值范围。然后看哪些层的范围特别大，对这些层单独设置精度。这样既保证了精度，又兼顾了性能。

常见坑四：内存爆炸

这个坑我踩过好几次。模型转换时，TensorRT会尝试做各种优化，比如算子融合、内存复用。但如果模型太大，或者workspace设置太小，就会爆内存。

解决办法：

  • 减小workspace：把 set_memory_pool_limit 设小一点，比如512MB。
  • 分步构建：先构建一个简化版，再逐步添加层。
  • 使用更小的batch size：构建时用batch=1，推理时再动态调整。

实战检查清单

每次做模型转换，我都会过一遍这个清单：

检查项 说明 工具/方法
ONNX结构检查 确保ONNX模型结构正确 onnx.checker.check_model()
算子兼容性 检查所有算子是否被TensorRT支持 polygraphy inspect model
精度对比 对比PyTorch和TensorRT的输出差异 自定义脚本，计算余弦相似度
性能测试 测试不同batch size下的推理延迟 trtexec --benchmarks
内存占用 检查推理时的显存占用 nvidia-smi

最后说两句

模型转换这件事，说白了就是个经验活。你踩的坑越多，后面就越顺。我刚开始做的时候，一个模型能折腾好几天。现在基本一两个小时就能搞定。

记住一点：不要相信任何一步的默认配置。PyTorch的默认导出参数不一定适合你，TensorRT的默认优化策略也不一定最优。每个模型都有自己的脾气，你得慢慢摸透它。

好了，这一章就聊到这儿。下一章咱们聊聊TensorRT的profile工具，看看怎么精准定位性能瓶颈。

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