Ultralytics YOLOv11姿态估计模型在WebGPU上的NMS导出问题解析

问题背景

在使用Ultralytics YOLOv11姿态估计模型时，开发者发现当在导出ONNX模型时启用NMS(nms=True)选项后，模型在WebGPU环境下运行时会出现错误，而在WASM(CPU)环境下则能正常工作。这一现象揭示了YOLO模型在不同推理后端上的兼容性差异。

错误分析

WebGPU运行时抛出的错误信息表明问题出在Pad操作节点上：

[E:onnxruntime:, sequential_executor.cc:516 ExecuteKernel] Non-zero status code returned while running Pad node. Name:'/Pad' Status Message: Failed to run JSEP kernel

Error: [WebGPU] Kernel "[Pad] /Pad" failed. Error: Program Pad has zero-sized tensor(s) in inputs or outputs. This is not supported now.

这类错误通常发生在张量形状处理环节，特别是当模型输出中存在动态维度或零尺寸张量时。WebGPU后端对这类情况的处理不如CPU后端灵活。

技术原理

YOLOv11姿态估计模型在导出为ONNX格式时，NMS操作会引入动态形状的输出。这是因为NMS处理后保留的检测框数量是不确定的，可能导致某些情况下输出张量的某些维度为零。在传统CPU环境中，ONNX运行时能够优雅地处理这种动态形状，但WebGPU后端目前对此支持有限。

解决方案

针对这一问题，项目维护者提出了一个代码修改方案：修改模型导出逻辑，确保即使NMS处理后保留的检测框数量为零，输出张量也能保持固定形状。具体修改是将原始代码中的动态赋值改为固定形状赋值：

原始代码： out[i] = dets

修改后： out[i, :keep.shape[0]] = dets

这一修改确保了输出张量始终保持预定义的形状，避免了WebGPU后端不支持的零尺寸张量情况。

实施建议

对于遇到类似问题的开发者，建议采取以下步骤：

1. 更新到最新版本的Ultralytics库，确保包含最新的修复和改进
2. 在导出模型时明确指定输出形状，避免动态维度
3. 如果必须使用WebGPU后端，考虑在模型后处理中添加形状检查逻辑
4. 对于关键应用，建议同时准备CPU和GPU两种推理方案以应对不同环境

总结

YOLOv11姿态估计模型在WebGPU上的NMS导出问题揭示了深度学习模型在不同推理后端上的兼容性挑战。通过理解底层技术原理和适当调整模型导出逻辑，开发者可以克服这些兼容性问题，实现模型在各种环境下的稳定运行。这一案例也提醒我们，在跨平台部署深度学习模型时，需要充分考虑不同推理后端的特点和限制。