YOLOv6模型导出ONNX时输入尺寸问题的解决方案

问题背景

在使用YOLOv6进行目标检测模型训练时，开发者经常会遇到需要将训练好的PyTorch模型导出为ONNX格式的需求。然而在实际操作中，当尝试导出使用特定输入尺寸训练的模型时，可能会遇到张量尺寸不匹配的错误。

错误现象分析

在YOLOv6项目中，当用户尝试使用非标准尺寸（如640x360）训练模型后，在导出ONNX模型时会出现以下典型错误：

RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 24 but got size 23 for tensor number 1 in the list.

这个错误发生在模型的前向传播过程中，具体是在特征图拼接时出现了尺寸不匹配的情况。

根本原因

YOLOv6网络架构设计中对输入图像尺寸有一个重要限制：输入尺寸必须是32的整数倍。这是因为：

1. YOLOv6网络包含多个下采样层，每层通常会将特征图尺寸减半
2. 经过5次下采样后，原始图像尺寸会被缩小32倍（2^5）
3. 如果原始尺寸不是32的倍数，会导致特征图尺寸出现小数，进而引发张量尺寸不匹配

解决方案

要解决这个问题，必须确保导出ONNX模型时指定的输入尺寸满足32的倍数要求。对于原问题中640x360的尺寸：

1. 高度360不是32的整数倍（360÷32=11.25）
2. 可以调整为最接近的32的倍数：352（32×11）或384（32×12）

因此正确的导出命令应为：

python ./deploy/ONNX/export_onnx.py --weights 640x360.pt --img 384 640 --batch 1 --simplify

或者：

python ./deploy/ONNX/export_onnx.py --weights 640x360.pt --img 352 640 --batch 1 --simplify

最佳实践建议

1. 训练时尺寸选择：建议在训练时就使用符合32倍数规则的输入尺寸，避免后续转换问题
2. 推理时尺寸处理：在实际部署中，如果输入图像尺寸不符合要求，应该先进行适当的填充或裁剪
3. 多尺寸验证：导出ONNX前，可以使用不同尺寸进行验证，确保模型鲁棒性
4. 性能考量：某些硬件设备对特定尺寸可能有优化，可针对性选择

总结

YOLOv6作为高效的目标检测框架，对输入尺寸有一定要求。理解网络架构中的下采样机制和尺寸约束，能够帮助开发者避免常见的模型导出问题。通过遵循32倍数的尺寸规则，可以确保模型从训练到部署的顺畅转换。