YOLOFuse如何导出ONNX模型？后续部署转换教程-编程阁

YOLOFuse如何导出ONNX模型？后续部署转换教程

在智能安防、自动驾驶和夜间监控等实际场景中，单一可见光图像的检测能力在低光照或复杂气象条件下常常捉襟见肘。面对这一挑战，多模态融合技术逐渐成为主流解决方案——尤其是RGB-红外（IR）双流检测，通过结合热成像与可见光信息，显著提升了系统在黑暗、烟雾或遮挡环境下的鲁棒性。

YOLO 系列以其高效性和实用性广受青睐，而基于 Ultralytics 架构构建的YOLOFuse框架，则专为双模态目标检测设计，在 LLVIP 等公开数据集上展现出优异性能。然而，训练完成的模型若要落地到边缘设备或异构硬件平台（如 Jetson、RK3588、华为 Ascend），必须跨越“从训练到推理”的鸿沟。此时，ONNX（Open Neural Network Exchange）作为跨平台模型交换的核心格式，便显得尤为重要。

掌握如何将 YOLOFuse 成功导出为 ONNX 并实现稳定推理，不仅是工程落地的关键一步，更是打通“训练—优化—部署”全链路的基础能力。

ONNX：为什么它是模型部署的“通用语言”？

ONNX 的本质是一个开放的神经网络中间表示（IR），它把不同框架中的模型统一为标准的计算图结构。PyTorch 训练的模型可以通过torch.onnx.export()转换为.onnx文件，之后可在 ONNX Runtime、TensorRT、OpenVINO 等多种推理引擎上运行，彻底摆脱对原始训练环境的依赖。

以一个典型的 YOLOv8 导出为例：

from ultralytics import YOLO model = YOLO("yolov8s.pt") model.export(format="onnx", imgsz=640, opset=13)

短短几行代码即可完成导出，得益于 Ultralytics 对主流架构的高度封装。但问题在于：YOLOFuse 是自定义双分支结构，无法直接调用.export()方法。它的输入是两个独立张量（RGB 和 IR 图像），融合逻辑也嵌入在网络中间，这些都会导致标准导出流程失败。

因此，我们需要深入理解其内部机制，并进行针对性适配。

YOLOFuse 的结构特性决定了导出方式

YOLOFuse 支持三种典型融合策略：

早期融合：将 RGB 与 IR 在通道维度拼接后送入同一主干网络（如[B,6,H,W]输入）。优点是特征交互充分，适合小目标；缺点是参数量翻倍。
中期融合：分别提取两路特征，在某个 stage 后融合（如 concat 或 add），再接入后续 neck 和 head。兼顾精度与效率，推荐用于资源受限场景。
决策级融合：两路完全独立推理，最终对 bbox 结果做 NMS 融合。鲁棒性强，但延迟高、内存占用大。

无论哪种方式，关键共性是：双输入、多分支、融合操作不可动态切换。

这意味着在导出 ONNX 时，必须确保：
1. 所有控制流（如if self.fusion_type == 'early'）已被固化；
2. 模型forward()接口明确接收(x_rgb, x_ir)两个输入；
3. 自定义模块（如注意力、融合层）使用的是 ONNX 支持的算子。

例如，一个简化版中期融合主干可能如下：

class DualBackbone(nn.Module): def __init__(self, backbone, fuse_at='stage2'): super().__init__() self.backbone_rgb = backbone() self.backbone_ir = backbone() self.fuse_conv = nn.Conv2d(512 * 2, 512, 1) # 假设输出通道为512 def forward(self, x_rgb, x_ir): f1_rgb = self.backbone_rgb.stage1(x_rgb) f1_ir = self.backbone_ir.stage1(x_ir) fused = torch.cat([f1_rgb, f1_ir], dim=1) fused = self.fuse_conv(fused) out = self.backbone_rgb.stage2(fused) return out

这类结构虽然清晰，但在导出时容易因 Python 控制流或非追踪兼容操作报错。解决思路只有一个：让整个前向过程可被 TorchScript 正确 trace 或 script。

如何安全导出 YOLOFuse 的 ONNX 模型？

以下是经过验证的完整导出流程，适用于大多数基于 PyTorch 的双流融合模型。

第一步：加载并冻结模型

务必在eval()模式下导出，关闭 Dropout、BatchNorm 统计更新等训练相关行为。

import torch from models.yolofuse import YOLOFuseNet # 替换为你自己的模型类 model = YOLOFuseNet(config="cfg/yolofuse_mid.yaml") model.load_state_dict(torch.load("weights/best_mid_fuse.pth")) model.eval() # 关键！

第二步：准备虚拟输入

注意输入形状需与训练一致，通常为(1, 3, 640, 640)单 batch 测试。

dummy_rgb = torch.randn(1, 3, 640, 640) dummy_ir = torch.randn(1, 3, 640, 640)

第三步：执行导出

torch.onnx.export( model, (dummy_rgb, dummy_ir), "yolofuse_mid.onnx", export_params=True, opset_version=13, do_constant_folding=True, input_names=["input_rgb", "input_ir"], output_names=["output"], dynamic_axes={ "input_rgb": {0: "batch_size", 2: "height", 3: "width"}, "input_ir": {0: "batch_size", 2: "height", 3: "width"}, "output": {0: "batch_size"} }, verbose=False )

几个关键参数说明：

opset_version=13：建议使用 11~13，避免旧版本不支持Resize、Where等常见算子；
dynamic_axes：允许变尺寸输入，提升部署灵活性；
input_names/output_names：命名便于后续绑定输入输出张量；
do_constant_folding：合并常量节点，减小模型体积。

⚠️ 若出现类似 “Can’t export operator aten::where” 错误，可能是 Opset 版本过低或算子未映射。可通过升级 PyTorch 和 ONNX，或重写子模块规避。

第四步：验证模型有效性

导出完成后一定要检查模型完整性：

import onnx onnx_model = onnx.load("yolofuse_mid.onnx") onnx.checker.check_model(onnx_model) print("✅ ONNX 模型验证通过")

这一步能捕获大部分结构错误，比如断连的节点、非法 shape 等。

部署实战：在 ONNX Runtime 上运行双流推理

一旦获得有效的.onnx文件，就可以在任意支持 ONNX 的平台上运行推理。以下是在 CPU/GPU 上使用 ONNX Runtime 的典型流程。

import onnxruntime as ort import numpy as np # 加载会话 sess = ort.InferenceSession( "yolofuse_mid.onnx", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] # 优先使用 GPU ) # 准备输入数据（假设已预处理） rgb_tensor = preprocess(rgb_image).numpy() # shape: (1,3,640,640) ir_tensor = preprocess(ir_image).numpy() # 推理 outputs = sess.run( None, { "input_rgb": rgb_tensor, "input_ir": ir_tensor } ) # outputs[0] 通常是 [batch, num_boxes, 85] 形式的原始输出

后处理部分仍需手动实现 bbox 解码、置信度筛选和 NMS，这部分逻辑与原 YOLO 保持一致。

实际部署中的工程考量

1. 融合策略选型建议

策略	推荐场景
中期融合	边缘设备部署首选，体积小、速度快，mAP 下降有限
早期融合	对小目标敏感任务，且算力充足
决策级融合	多传感器异构系统，容忍较高延迟

根据实测数据，中期融合在仅增加 2.61MB 模型体积的情况下达到 94.7% mAP@50，性价比极高。

2. 输入同步性保障

RGB 与 IR 图像必须严格时间对齐。若存在帧偏移，会导致目标位置错位，严重影响融合效果。建议：
- 使用硬件触发信号同步采集；
- 或在软件层添加时间戳匹配机制。

3. 异常处理与降级策略

真实环境中可能出现单路图像丢失（如 IR 模块故障）。理想做法是设计 fallback 机制：

if ir_signal_lost: # 自动退化为 RGB-only 模式 output = single_stream_model(rgb_input) else: output = dual_stream_model(rgb_input, ir_input)

这种容错能力极大提升系统稳定性。

4. 性能优化路径

导出后的 ONNX 模型还可进一步优化：

开启 ONNX Runtime 图优化：
python sess_options = ort.SessionOptions() sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess = ort.InferenceSession("yolofuse_mid.onnx", sess_options, ...)
转 TensorRT 引擎加速：
利用onnx-tensorrt工具链可将 ONNX 编译为高效 TRT 引擎，在 NVIDIA 设备上获得 3~5 倍速度提升。
量化压缩：
支持 FP16/INT8 量化以降低显存占用和功耗，特别适合移动端部署。