TorchMetrics中多标签混淆矩阵绘图异常分析与修复

问题背景

在使用TorchMetrics库的MultilabelConfusionMatrix功能时，当处理两个标签的分类任务时，绘图功能会出现异常。具体表现为调用plot()方法时抛出AttributeError，提示'numpy.ndarray'对象没有'set_title'属性。

问题现象

当用户尝试为两个标签绘制混淆矩阵时，系统会在torchmetrics/utilities/plot.py文件的第275行报错。错误信息表明程序试图在一个numpy数组上调用set_title方法，而实际上这个数组应该是一个matplotlib的Axes对象。

问题分析

通过分析源代码，我们发现问题的根源在于子图(Subplot)处理逻辑的条件判断不当。在当前的实现中，只有当行数和列数都不等于1时才会获取子图数组中的单个Axes对象，否则直接使用整个Axes数组。这种判断方式在处理2×1或1×2的子图布局时会出现问题。

技术细节

在matplotlib中，当创建多个子图时：

• 单子图情况：直接返回一个Axes对象
• 多行多列子图：返回一个Axes对象数组
• 单行多列或多行单列：也返回一个Axes对象数组

当前的判断条件rows != 1 and cols != 1只处理了多行多列的情况，忽略了单行多列或多行单列的情况，导致程序错误地将数组当作单个Axes对象使用。

解决方案

将条件判断从and改为or：

ax = axs[i] if rows != 1 or cols != 1 else axs

这样修改后，只要行数或列数中有一个不等于1，就会正确地获取子图数组中的单个Axes对象，只有在真正的单子图情况下才会直接使用整个Axes对象。

影响范围

该问题影响所有使用TorchMetrics库中MultilabelConfusionMatrix功能且标签数为2的用户。标签数大于2的情况不受影响，因为此时会创建多行多列的子图布局；标签数为1的情况已经被正确处理。

修复验证

修复后，用户可以正常地为两个标签的分类任务绘制混淆矩阵，每个标签的混淆矩阵会正确地显示在独立的子图中，并且每个子图都会有正确的标题标注。

总结

这个案例展示了在开发通用绘图工具时需要考虑各种可能的子图布局情况。特别是在处理动态生成的图表时，必须全面考虑所有可能的输入参数组合，确保程序在各种情况下都能正确工作。