二值语义分割自定义Weighted IoU Loss实现正确性问询

我正在基于卫星图像进行道路与前景像素的二值语义分割任务，想咨询我实现的加权IoU损失（Weighted IoU Loss）是否正确。

我的U-Net模型最后一层是1x1卷积层，之后经过sigmoid激活，输出形状为(batch_size, 400, 400, 1)，每个像素值代表该像素属于标记为1的类别的概率。批量大小设置为8。添加类别权重是因为前景像素与道路像素的比例为9:1，但训练时损失在几个epoch后就停滞不前，调整学习率也没解决问题。

以下是我实现的加权IoU损失计算代码：

def compute_loss(self, y_true, y_pred):
    
    '''Implementation code to calculate weighted IoU loss for U-Net model. 
    
    Args:
        y_true: tensor of shape (batch_size, 400, 400, 1), containing pixel values of groundtruth image
        y_pred: tensor of shape (batch_size, 400, 400, 1), containing pixel values of prediction from model

    Returns:
        loss
    '''        
    
    scores = []
    class_weights = [self.foreground_weight, self.road_weight]
    
    for i in range(y_true.shape[0]):
        
        y_true_sample = tf.reshape(y_true[i], [-1])
        y_pred_sample = tf.reshape(y_pred[i], [-1])
        
        intersection = tf.reduce_sum(tf.cast(y_true_sample, tf.float32)*tf.cast(y_pred_sample, tf.float32))
        smoothing = 1
        iou = (intersection + smoothing)/ (tf.reduce_sum(tf.cast(y_true_sample, tf.float32)) + tf.reduce_sum(tf.cast(y_pred_sample, tf.float32)) - intersection + smoothing)
        
        weighted_iou = tf.reduce_sum((1. - iou)*class_weights)
        scores.append(weighted_iou)
        
    scores = tf.convert_to_tensor(scores, dtype = tf.float32)
    return tf.reduce_mean(scores)

• 当前实现的核心问题
  你的加权IoU实现完全没有起到“加权”作用。当前代码中，每个样本计算出单个IoU后，用(1 - iou)乘以类别权重列表再求和，本质只是给(1 - iou)乘了一个固定系数（前景权重+道路权重），完全没有针对类别不平衡做针对性加权，和普通IoU损失的效果没有区别。

• 二值分割加权IoU的正确思路
  二值分割场景下，加权IoU需要针对道路（正样本）和前景（负样本）的交、并计算分别引入权重，或者直接给不同类别的像素赋予权重后再计算整体IoU。常用的两种方式：

  1. 分别计算两类的IoU，乘以对应权重后求和
  2. 在交、并的像素求和阶段，给每个像素乘以所属类别的权重

• 修正后的代码示例
  下面是适配你场景的加权IoU损失实现，直接在交并计算中引入类别权重，同时去掉低效的batch循环：

  def compute_loss(self, y_true, y_pred):
      '''加权IoU损失（适配二值语义分割）
      
      Args:
          y_true: 形状(batch_size, 400, 400, 1)的真实标签张量
          y_pred: 形状(batch_size, 400, 400, 1)的模型预测概率张量（sigmoid输出）
  
      Returns:
          平均损失值
      '''
      # 统一数据类型
      y_true = tf.cast(y_true, tf.float32)
      y_pred = tf.cast(y_pred, tf.float32)
      
      # 类别权重：道路（y=1）占比10%，赋予更高权重平衡类别
      road_weight = self.road_weight  # 建议设为9
      foreground_weight = self.foreground_weight  # 建议设为1
      
      # 计算加权交集：道路像素的交集乘道路权重，前景像素的交集乘前景权重
      intersection = tf.reduce_sum(y_true * y_pred * road_weight) + tf.reduce_sum((1 - y_true) * (1 - y_pred) * foreground_weight)
      
      # 计算加权并集：避免重复计算交集
      union = (tf.reduce_sum(y_true * road_weight) + tf.reduce_sum(y_pred * road_weight) +
               tf.reduce_sum((1 - y_true) * foreground_weight) + tf.reduce_sum((1 - y_pred) * foreground_weight) -
               intersection)
      
      # 用小平滑值避免除零，同时不干扰真实IoU计算
      smoothing = 1e-6
      iou = (intersection + smoothing) / (union + smoothing)
      
      # 损失为1减去加权IoU
      return 1. - iou
  

• 额外优化建议

  • 去掉遍历batch的循环，用TensorFlow批量运算提升效率，同时避免张量转换的潜在问题
  • 平滑值改用1e-6而非1，过大的平滑值会掩盖IoU的真实变化，阻碍模型学习
  • 若单纯加权IoU效果不佳，可尝试混合损失：加权二元交叉熵 + IoU损失，兼顾像素级分类精度和整体交并比优化
  • 检查数据预处理：确保标签中道路像素为1、前景为0，且数据增强没有破坏类别比例的合理性

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