基于Keras的U-Net二分类语义分割训练时MeanIoU不上升问题求解

根因分析 #

• 核心问题是tf.keras.metrics.MeanIoU要求输入为离散的整数类别标签，不会自动处理sigmoid输出的0~1连续概率值：
  你的模型输出层使用sigmoid激活，输出是0到1之间的概率结果。Keras自带的二分类准确率binary_accuracy默认内置了0.5的阈值逻辑，会自动将概率值转换为0/1的二值类别后再计算准确率，因此训练过程中准确率会随着损失下降正常上升；但MeanIoU没有内置二值化逻辑，会直接将输入的浮点概率值当做类别索引处理，无法正确统计混淆矩阵，因此指标一直卡在固定值没有变化。
• 次要可能原因：如果数据集存在严重类别不平衡（比如背景占比远高于前景），也可能出现仅靠预测大部分背景就能拉高准确率，但前景预测效果没有提升、IoU不上涨的情况，结合你给出的训练日志损失下降幅度较大的特征，大概率为上述第一个原因。

解决方法 #

自定义带二值化逻辑的MeanIoU指标，在计算前先对模型输出的概率值做阈值截断，再传入原生MeanIoU计算即可，具体代码修改如下：

1. 首先自定义二分类MeanIoU指标

import tensorflow as tf

class BinaryMeanIoU(tf.keras.metrics.MeanIoU):
    def __init__(self, num_classes=2, threshold=0.5, name='binary_mean_iou', **kwargs):
        super().__init__(num_classes=num_classes, name=name, **kwargs)
        self.threshold = threshold

    def update_state(self, y_true, y_pred, sample_weight=None):
        # 对预测概率做二值化处理
        y_pred = tf.cast(tf.greater_equal(y_pred, self.threshold), tf.int32)
        y_true = tf.cast(y_true, tf.int32)
        return super().update_state(y_true, y_pred, sample_weight)

2. 替换模型编译时的原有MeanIoU指标

UNet.compile(
    optimizer="Adam",
    loss="binary_crossentropy", 
    metrics=["accuracy", BinaryMeanIoU(num_classes=2)]
)

如果确认数据集存在类别不平衡问题，可以同时做以下优化：

• 替换损失函数为Dice损失、Focal Loss等对类别不平衡更友好的损失，或者为交叉熵设置类别权重
• 根据验证集的前景占比调整二值化阈值，不要固定使用0.5，进一步提升IoU指标

内容的提问来源于stack exchange，提问作者Liana Olianov