📚 课程目录第4章 YOLO模型导出

4. YOLO模型导出：从PyTorch到ONNX再到TensorRT的模型转换全流程

模型训练好只是第一步。真正要让它跑在嵌入式设备上，还得过一道关——模型导出。说白了，就是把PyTorch的权重文件，变成TensorRT能吃的引擎文件。这个过程我踩过不少坑，今天一次性讲清楚。

核心链路：PyTorch模型 → ONNX中间表示 → TensorRT引擎

4.1 为什么不能直接用PyTorch跑推理？

你想想看，PyTorch设计出来是为了训练。它里面有很多动态图、自动求导的机制，推理时根本用不上。而且PyTorch对INT8量化、层融合、内存复用这些优化手段支持得很有限。

TensorRT就不一样了。它专门为推理做了极致优化。我做过对比，同一个YOLOv5s模型，PyTorch直接推理要跑35ms，转成TensorRT FP16后只要12ms。差距就是这么明显。

推理后端	精度	延迟（ms）	适用场景
PyTorch	FP32	35	训练、验证
ONNX Runtime	FP32	28	跨平台部署
TensorRT FP16	FP16	12	嵌入式推理
TensorRT INT8	INT8	8	极致性能

4.2 第一步：PyTorch导出ONNX

ONNX就像一个中间人。它不挑框架，PyTorch、TensorFlow都能转成它。我个人习惯先用ONNX做一次验证，确保模型结构没问题。

导出代码其实不复杂，但有几个关键点要注意。

import torch
import torch.onnx

# 加载训练好的模型
model = torch.load('yolov5s.pt', map_location='cpu')['model'].float()
model.eval()

# 构造一个虚拟输入
dummy_input = torch.randn(1, 3, 640, 640)

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    'yolov5s.onnx',
    opset_version=11,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

我的经验：opset_version别用太高的，11或12最稳。我之前试过13，结果某些算子不兼容，折腾了一整天。

导出后一定要验证一下。用ONNX Runtime跑一遍，看输出和PyTorch是否一致。

import onnx
import onnxruntime as ort

# 检查模型结构
onnx_model = onnx.load('yolov5s.onnx')
onnx.checker.check_model(onnx_model)

# 用ONNX Runtime推理
session = ort.InferenceSession('yolov5s.onnx')
outputs = session.run(None, {'input': dummy_input.numpy()})

4.3 第二步：ONNX转TensorRT

这一步是重头戏。TensorRT提供了两种转换方式：trtexec命令行工具和Python API。我建议新手先用命令行，参数清晰，不容易出错。

4.3.1 用trtexec转换

trtexec --onnx=yolov5s.onnx \
        --saveEngine=yolov5s_fp16.engine \
        --fp16 \
        --workspace=1024 \
        --minShapes=input:1x3x640x640 \
        --optShapes=input:4x3x640x640 \
        --maxShapes=input:8x3x640x640

这里有个坑——动态shape。如果你的应用场景需要不同batch size，一定要设置--minShapes、--optShapes和--maxShapes。我曾经没设这个，结果部署时输入尺寸一变就报错。

4.3.2 用Python API转换

如果你需要更精细的控制，比如自定义插件，那就用Python API。

import tensorrt as trt

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

# 解析ONNX
with open('yolov5s.onnx', 'rb') as f:
    parser.parse(f.read())

# 配置builder
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
config.set_flag(trt.BuilderFlag.FP16)

# 构建引擎
serialized_engine = builder.build_serialized_network(network, config)
with open('yolov5s_fp16.engine', 'wb') as f:
    f.write(serialized_engine)

注意：TensorRT版本要和CUDA、cuDNN匹配。我见过有人用TRT 8.5配CUDA 11.0，结果编译都过不去。建议用TRT官方推荐的版本组合。

4.4 整个流程的SVG图

下面这张图把整个流程串起来了。你可以对照着看，每一步做什么，输出什么。

4.5 常见问题与避坑指南

转换过程中问题不少。我挑几个典型的说说。

算子不支持：YOLO里有些自定义算子，TensorRT不认识。解决办法是写插件，或者用ONNX的custom_op替换。我曾经花了两天写一个SiLU激活函数的插件。
动态shape报错：如果输入尺寸会变，一定要设dynamic_axes和optShapes。否则推理时输入尺寸不对，直接崩。
精度下降：FP16一般没问题，INT8就需要校准数据集了。我建议先用FP16，性能不够再上INT8。
内存不足：嵌入式设备显存有限。转换时workspace别设太大，否则直接OOM。

我的习惯：每次转换完，我都会用同一张图片跑一遍PyTorch和TensorRT，对比输出结果的mAP。偏差在1%以内就算合格。

4.6 总结一下

模型导出这件事，说难不难，说简单也不简单。核心就是三步：PyTorch转ONNX，ONNX转TensorRT，验证一致性。每一步都有坑，但只要你按流程来，多测试，基本不会出大问题。