MMDeploy中RetinaNet单类别模型导出问题的分析与解决

问题背景

在使用MMDeploy工具导出基于MMDetection的RetinaNet模型时，当模型配置为单类别检测任务且使用CrossEntropyLoss作为分类损失函数时，会出现导出失败的问题。具体表现为在模型导出过程中抛出"IndexError: max(): Expected reduction dim 2 to have non-zero size"错误。

技术分析

RetinaNet模型结构特点

RetinaNet是一种经典的单阶段目标检测器，其核心由特征金字塔网络(FPN)和两个子网络(分类子网络和回归子网络)组成。在MMDetection实现中，RetinaHead负责处理这两个子网络的输出。

问题根源

该问题源于MMDeploy中对RetinaNet模型导出时的特殊处理逻辑。具体来说，当配置为：

1. 单类别检测(num_classes=1)
2. 使用CrossEntropyLoss作为分类损失(use_sigmoid=False)

时，MMDeploy中的base_dense_head.py文件对分类分数进行了两次切片操作，导致最终用于非极大值抑制(NMS)的分数张量维度变为0，从而引发错误。

代码逻辑分析

在MMDeploy的base_dense_head.py中，处理流程如下：

1. 第一次切片：从原始分类分数中排除背景类别(当use_sigmoid=False时)
2. 第二次切片：再次从分数中排除最后一个类别(当use_sigmoid=False时)

对于单类别情况，经过这两次切片后，分类分数张量的最后一个维度变为0，导致后续的max()操作失败。

解决方案

临时解决方案

对于遇到此问题的用户，可以采取以下临时解决方案：

1. 修改模型配置，使用Sigmoid激活函数(use_sigmoid=True)
2. 手动修改MMDeploy源代码，移除其中一次不必要的切片操作

根本解决方案

该问题的根本解决方案是优化MMDeploy中对RetinaNet模型的导出逻辑：

1. 对于单类别情况，应避免重复切片操作
2. 统一处理不同类别数量情况下的分数处理逻辑
3. 增加对边界条件的检查和处理

最佳实践建议

1. 对于单类别检测任务，建议优先考虑使用Sigmoid激活函数
2. 导出模型前，建议先验证模型在原生框架中的推理功能
3. 关注MMDeploy的版本更新，及时获取官方修复

总结

MMDeploy作为模型部署工具，在支持各种检测模型导出时需要考虑多种边界条件。RetinaNet单类别导出问题揭示了在模型转换过程中对特殊配置情况的处理不足。通过理解问题本质和解决方案，用户可以更好地完成模型部署工作，同时也为开发者提供了优化工具的方向。