3、模型导出与转换：PyTorch模型转ONNX、ONNX转TensorRT引擎、动态batch支持、精度校准（FP16/INT8）

好，咱们进入实战环节的第一个硬骨头——模型导出与转换。

说实话，很多同学在训练完一个YOLO模型后，觉得万事大吉了。结果一到部署，发现推理速度慢得离谱，或者干脆跑不起来。为什么？因为PyTorch模型是给训练用的，不是给生产环境用的。它太重了，依赖太多。

我个人习惯，训练完模型第一件事就是把它转成ONNX。ONNX就像一个中间语言，PyTorch、TensorFlow都能转过去，然后从ONNX再转到各个平台的专属格式。今天咱们就一步步来，把YOLOv5/v8从PyTorch一路转到TensorRT引擎，顺便把动态batch和精度校准也搞定。

3.1 PyTorch模型转ONNX

这一步其实不复杂，但有几个坑。我先说标准做法。

YOLOv5官方仓库里其实已经提供了导出脚本，但咱们得理解它干了什么，不然出了问题你都不知道怎么修。

import torch
import torch.onnx

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

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

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

这里有个关键点——opset_version。我建议用12或更高版本。为什么？因为低版本对某些算子支持不好，尤其是YOLO里用的SiLU激活函数和Focus层。我曾经在opset 11上折腾了一整天，导出倒是成功了，但推理结果全是乱的。后来换成opset 12，问题直接消失。

注意：导出时一定要用 model.eval() 模式。如果你忘了，模型里的BatchNorm和Dropout会保持训练状态，导出的ONNX推理结果会完全不对。我见过有人因为这个排查了三天。

3.2 ONNX转TensorRT引擎

拿到ONNX文件之后，下一步就是用TensorRT把它编译成引擎文件（.engine或.trt）。这一步才是真正性能起飞的地方。

TensorRT会做几件事：层融合、内存优化、精度校准。说白了，就是把你的计算图重新梳理一遍，能合并的合并，能去掉的去掉。

我一般用 trtexec 命令行工具来做转换，简单直接：

trtexec --onnx=yolov5s.onnx \
        --saveEngine=yolov5s_fp16.engine \
        --fp16 \
        --workspace=4096 \
        --minShapes=images:1x3x640x640 \
        --optShapes=images:8x3x640x640 \
        --maxShapes=images:16x3x640x640

看到 --minShapes、--optShapes、--maxShapes 了吗？这就是动态batch的支持方式。你告诉TensorRT：我最小batch是1，最优batch是8，最大batch是16。TensorRT会针对这些尺寸分别做优化，运行时根据实际输入自动选择最优的推理路径。

小技巧：如果你的应用场景batch大小是固定的，比如永远是1，那就别开动态batch。把 --minShapes 和 --maxShapes 设成一样，TensorRT能做得更极致的优化。我做过对比，固定batch比动态batch能快5%-10%。

3.3 动态batch支持

刚才提到了动态batch，这里展开说说。

动态batch说白了就是允许你在推理时每次传入不同数量的图片。比如这次来1张，下次来8张，引擎都能处理。这在线上服务场景特别有用——请求量波动大，你不能每次都凑够一个固定batch再推理。

在代码里怎么用？以Python API为例：

import tensorrt as trt
import pycuda.driver as cuda

# 加载引擎
with open('yolov5s_dynamic.engine', 'rb') as f:
    engine_data = f.read()
runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
engine = runtime.deserialize_cuda_engine(engine_data)

# 创建上下文
context = engine.create_execution_context()

# 动态设置输入形状
input_shape = (batch_size, 3, 640, 640)
context.set_binding_shape(0, input_shape)

# 分配内存
input_size = trt.volume(input_shape) * trt.float32.itemsize
# ... 后续推理代码

关键就一行：context.set_binding_shape(0, input_shape)。每次推理前，根据实际batch大小设置一下就行。注意，这个batch大小不能超过你导出时指定的 --maxShapes。

3.4 精度校准（FP16/INT8）

精度校准是TensorRT最核心的功能之一。说白了就是用更少的数据位来表示模型参数和中间结果，从而大幅提升速度、降低显存占用。

FP16 比较简单，几乎是无损的。YOLO这种模型，FP16推理和FP32推理的mAP差距通常在0.1%以内，肉眼根本看不出来。我建议所有部署场景都开FP16，白捡的性能提升。

INT8 就有点讲究了。它需要校准数据集，因为模型要学习如何用8位整数近似表示原来的浮点数。校准集一般选500-1000张有代表性的图片，覆盖你实际场景中的各种情况。

TensorRT提供了INT8校准的接口，我一般用 trtexec 配合校准图片来做：

trtexec --onnx=yolov5s.onnx \
        --saveEngine=yolov5s_int8.engine \
        --int8 \
        --calib=/path/to/calibration_images/ \
        --workspace=4096

如果你用Python API，需要自己写一个校准器类：

class YOLOCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calibration_files, batch_size, input_shape):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.calibration_files = calibration_files
        self.batch_size = batch_size
        self.input_shape = input_shape
        self.batch_idx = 0
        # 分配GPU内存
        self.d_input = cuda.mem_alloc(trt.volume(input_shape) * trt.float32.itemsize)

    def get_batch(self, names):
        if self.batch_idx * self.batch_size >= len(self.calibration_files):
            return None
        # 读取并预处理图片
        batch = self._load_batch()
        cuda.memcpy_htod(self.d_input, batch)
        self.batch_idx += 1
        return [int(self.d_input)]

    def get_batch_size(self):
        return self.batch_size

    def read_calibration_cache(self):
        return None

    def write_calibration_cache(self, cache):
        with open('calibration.cache', 'wb') as f:
            f.write(cache)
经验之谈：INT8校准后，YOLOv5s的mAP下降通常在0.5%-1%之间，但推理速度能提升2-3倍。如果你对精度要求极高（比如自动驾驶），建议先用FP16。如果追求极致性能（比如边缘设备），INT8是必选项。我曾在Jetson Nano上跑YOLOv5s，FP16只能跑15FPS，INT8直接干到30FPS，翻倍。

3.5 避坑指南

最后，分享几个我踩过的坑：

  • ONNX导出时算子不支持：YOLOv5的Focus层在旧版ONNX里表现不好。解决方案是升级opset版本，或者在导出前手动把Focus层替换成普通的Conv层。我建议直接升级opset。
  • TensorRT编译时显存不足：--workspace 设太大，显卡扛不住。我一般设4096（4GB），如果报错就降到2048。
  • INT8校准后精度暴跌：大概率是校准集选得不好。比如你拿全是白天场景的图片去校准，结果部署时遇到夜晚场景，精度肯定崩。校准集一定要覆盖所有光照、角度、目标大小的情况。
  • 动态batch推理时形状不匹配：检查一下你的后处理代码，是不是写死了batch size。我见过有人引擎支持动态batch，但后处理里写死了 for i in range(1)，结果batch大于1时直接报错。

嗯，模型导出与转换这块，说白了就是「训练完别急着高兴，部署才是真功夫」。把ONNX和TensorRT玩熟了，你的YOLO模型才能真正从实验室走向生产线。下一章咱们聊聊推理引擎的搭建和优化，到时候你会看到，前面这些转换工作到底值不值。

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