Keras中UNet网络如何使用自定义权重图优化边界分割效果

你需要将预计算的逐像素权重图传入损失函数，你当前使用的全局类加权二元交叉熵仅能处理类别不平衡问题，无法覆盖UNet论文中针对相邻目标边界的动态权重优化需求，边界区域的高权重值需要逐像素乘入对应位置的损失计算，才能实现更优的边界分离效果。

推荐实现方案（适配Keras框架，无需修改模型结构） #

1. 调整训练标签格式 #

在数据预处理阶段，将同尺寸的掩码（mask）与预计算的逐像素权重图在最后一个通道维度拼接，原来的单通道掩码shape为(image_height, image_width, 1)，拼接后得到的新标签shape为(image_height, image_width, 2)，其中[..., 0:1]为原二值分割掩码，[..., 1:2]为逐像素权重图。

如果使用自定义数据生成器，将原来返回(image, mask)的逻辑修改为返回(image, np.concatenate([mask, weight_map], axis=-1))即可。

2. 重写自定义损失函数 #

新的损失函数会先拆分拼接的标签，再将逐像素权重乘入交叉熵计算：

def unet_weighted_binary_crossentropy(y_true, y_pred):
    # 拆分掩码和逐像素权重图
    mask = y_true[..., 0:1]
    pixel_weight = y_true[..., 1:2]
    # 数值截断避免对数计算溢出
    y_pred = K.clip(y_pred, K.epsilon(), 1 - K.epsilon())
    # 计算逐像素交叉熵后乘以对应权重
    cross_entropy = -(mask * K.log(y_pred) + (1 - mask) * K.log(1 - y_pred))
    weighted_loss = cross_entropy * pixel_weight
    # 返回平均损失值
    return K.mean(weighted_loss)

你之前unet_weight_map函数里已经将类权重wc和边界权重合并到了输出的权重图中，无需再向损失函数单独传入全局类权重数组。

3. 编译训练模型 #

模型编译时直接传入新的损失函数即可：

model.compile(optimizer='adam', loss=unet_weighted_binary_crossentropy, metrics=['accuracy'])

其他可选实现方式 #

你也可以使用Keras原生的逐像素样本加权能力实现：

• 编译模型时指定sample_weight_mode="pixelwise"
• 训练时在model.fit的sample_weight参数传入预计算的逐像素权重数组，数组shape为(样本数, image_height, image_width)
• 损失函数可直接使用标准二元交叉熵，Keras会自动将逐像素权重乘入损失计算。该方案适合使用tf.data构建训练流水线的场景，不需要修改标签的通道结构。

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