【YOLO 系列模型从 ONNX 到 TensorRT Engine 的转换指南：关键步骤与常见问题】

将 YOLO 系列模型的 .onnx 文件导出为 TensorRT 的 .engine 文件时，需要注意以下关键步骤和可能的修改点：

---

1. ONNX 导出前的模型调整

在导出 ONNX 模型前，需确保 YOLO 模型的某些操作与 TensorRT 兼容：

• 替换不支持的算子：
  • 例如，YOLOv5 的 Focus 层需要替换为等效的 Conv 层（YOLOv5 官方代码已支持此修改）。
  • 上采样（Upsample）操作应使用 Resize 而非 Upsample（确保插值模式为 nearest）。
• 静态/动态输入尺寸：
  • 如果模型需要支持动态输入（如不同分辨率），需在导出 ONNX 时设置 dynamic_axes。例如：

    torch.onnx.export(
        model,
        im,
        "yolo.onnx",
        dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}},
        # ...
    )
    

  • 静态输入则需固定尺寸（如 640x640）。
• 简化后处理：
  • YOLO 原生后处理（如非极大值抑制 NMS）可能包含 TensorRT 不支持的算子。建议导出时仅保留原始输出层，将后处理移至模型外部。

---

2. TensorRT 转换关键步骤

使用 TensorRT 的 trtexec 或 Python API 转换时需注意：

(1) 使用 trtexec 命令行工具

trtexec --onnx=yolo.onnx \
        --saveEngine=yolo.engine \
        --fp16 \                # 启用 FP16 加速（可选）
        --workspace=4096 \       # 分配显存空间
        --minShapes=images:1x3x640x640 \    # 动态输入的最小尺寸
        --optShapes=images:1x3x640x640 \     # 最优尺寸
        --maxShapes=images:1x3x640x640       # 动态输入的最大尺寸

• 动态输入：若需支持动态输入，需设置 minShapes/optShapes/maxShapes。
• 精度：添加 --fp16 或 --int8 提升速度（需硬件支持）。

(2) Python API 转换

使用 TensorRT 的 Python API 可更灵活控制：

import tensorrt as trt

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

with open("yolo.onnx", "rb") as f:
    parser.parse(f.read())

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

serialized_engine = builder.build_serialized_network(network, config)
with open("yolo.engine", "wb") as f:
    f.write(serialized_engine)

---

3. 常见问题与修复

(1) 算子不支持

• 错误示例：Unsupported ONNX operator: NonMaxSuppression
  • 修复：导出 ONNX 时不要包含 NMS 层，使用 TensorRT 的 EfficientNMS 插件或自定义后处理。
• 插件需求：
  • 对于某些算子（如 SPPF），需编译 TensorRT 插件（参考 TensorRT 官方插件库）。
  • 加载插件示例：

    trtexec --plugins=path/to/libcustom_plugin.so ... 
    

(2) 输入/输出维度不匹配

• 检查输入名称和形状：
  • 使用 Netron 可视化 ONNX 模型，确认输入节点名称（如 images:0）和维度。
  • 转换时需指定正确的输入名称和维度。

(3) 动态轴设置错误

• 错误示例：Invalid dynamic dimensions for input tensor
  • 修复：确保 dynamic_axes 在 ONNX 导出时正确设置，并在转换时指定 minShapes/optShapes/maxShapes。

---

4. 后处理集成

• 方案 1：分离后处理
  • ONNX 模型仅输出原始检测层（如 output:0），在 TensorRT 推理后自行实现 NMS。
• 方案 2：集成 NMS 插件
  • 使用 TensorRT 的 EfficientNMS 插件（需 YOLO 模型支持）。

---

5. 验证结果

• 一致性检查：

  import onnxruntime as ort
  import tensorrt as trt
  
  # 对比 ONNX 和 TensorRT 的输出是否一致
  ort_session = ort.InferenceSession("yolo.onnx")
  trt_runtime = trt.Runtime(trt.Logger())
  with open("yolo.engine", "rb") as f:
      engine = trt_runtime.deserialize_cuda_engine(f.read())
  

• 性能测试：使用 trtexec --avgRuns=100 测试推理速度。

---

总结

关键修改点包括：调整 ONNX 模型结构、处理动态输入、修复不支持的算子、集成后处理插件。根据具体错误信息调整转换参数或模型结构。