问题背景 #

我用Keras实现了U-Net模型，用来对显微镜图像中的细胞器做分割（参考U-Net论文）。为了让网络能识别出仅由1个像素分隔的多个独立目标，我打算给每个标签图像使用权重图（权重计算公式参考上述论文）。

我知道需要自定义交叉熵损失函数来利用这些权重图，但Keras的自定义损失函数默认只接受y_true和y_pred两个参数，不知道怎么把权重图的值加入到损失计算里。另外，我想问问能不能把权重图的值和标签值合并到y_true张量中？我是深度学习新手，要是问题表述有问题请见谅，恳请各位大佬给点帮助和建议！

我目前的自定义损失函数代码如下：

def pixelwise_crossentropy(self, ytrue, ypred):
    ypred /= tf.reduce_sum(ypred, axis=len(ypred.get_shape()) - 1, keep_dims=True)
    # manual computation of crossentropy
    _epsilon = tf.convert_to_tensor(epsilon, ypred.dtype.base_dtype)
    output = tf.clip_by_value(ypred, _epsilon, 1. - _epsilon)
    return - tf.reduce_sum(ytrue * tf.log(output))

解决方案 #

作为过来人，给你几个可行的思路，按从易到难排序：

方案1：把权重图和标签合并到y_true张量中

这是你想到的思路，完全可行，也是最容易实现的。具体来说，假设你的原始分割标签是单通道的掩码图，权重图也是同尺寸的单通道图，你可以把它们在最后一个维度拼接成一个2通道的张量，作为模型训练时的y_true输入。

修改后的损失函数可以这样写：

def weighted_pixelwise_crossentropy(ytrue, ypred):
    # 从ytrue中拆分出标签和权重图
    # 假设第0通道是原始标签，第1通道是权重图
    y_labels = ytrue[..., 0:1]
    weights = ytrue[..., 1:2]
    
    # 保留你原来的交叉熵计算逻辑
    ypred /= tf.reduce_sum(ypred, axis=len(ypred.get_shape()) - 1, keep_dims=True)
    # 建议把epsilon设成固定小值，比如1e-7，避免外部变量依赖
    _epsilon = tf.convert_to_tensor(1e-7, ypred.dtype.base_dtype)
    output = tf.clip_by_value(ypred, _epsilon, 1. - _epsilon)
    cross_entropy = - y_labels * tf.log(output)
    
    # 应用权重图计算最终损失
    weighted_loss = tf.reduce_sum(weights * cross_entropy)
    return weighted_loss

使用方式：训练前，把你的标签数组和权重数组在最后一个维度拼接，比如用NumPy：

# 假设labels是形状为(样本数, 高, 宽, 1)的标签张量
# weights是形状为(样本数, 高, 宽, 1)的权重张量
combined_y_true = np.concatenate([labels, weights], axis=-1)

之后直接把combined_y_true作为模型训练时的目标输入即可。

方案2：用闭包传递权重图（适合固定权重或全局权重）

如果不想修改y_true的结构，可以用Python闭包让损失函数访问外部的权重张量。比如：

def get_weighted_crossentropy(weights_tensor):
    def pixelwise_crossentropy(ytrue, ypred):
        ypred /= tf.reduce_sum(ypred, axis=len(ypred.get_shape()) - 1, keep_dims=True)
        _epsilon = tf.convert_to_tensor(1e-7, ypred.dtype.base_dtype)
        output = tf.clip_by_value(ypred, _epsilon, 1. - _epsilon)
        cross_entropy = - ytrue * tf.log(output)
        # 应用外部传入的权重
        return tf.reduce_sum(weights_tensor * cross_entropy)
    return pixelwise_crossentropy

不过这个方案更适合全局固定权重的场景，如果每个样本都有不同的权重图，你需要把权重图作为模型的第二个输入，然后结合自定义层或者训练循环来处理，相对麻烦一些。

方案3：自定义训练循环（最灵活，适合复杂场景）

如果你用的是TensorFlow 2.x版本，自定义训练循环会给你最大的灵活性——你可以直接把图像、标签、权重图作为三个独立的输入，在训练步骤里直接计算加权损失：

# 假设已经定义好你的U-Net模型
model = build_unet_model()
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)

@tf.function
def train_step(images, labels, weights):
    with tf.GradientTape() as tape:
        preds = model(images, training=True)
        # 计算交叉熵
        preds /= tf.reduce_sum(preds, axis=-1, keep_dims=True)
        _epsilon = tf.convert_to_tensor(1e-7, preds.dtype.base_dtype)
        clipped_preds = tf.clip_by_value(preds, _epsilon, 1. - _epsilon)
        cross_entropy = - labels * tf.log(clipped_preds)
        # 应用权重，这里用reduce_mean还是reduce_sum看你的需求
        weighted_loss = tf.reduce_mean(weights * cross_entropy)
    
    # 计算梯度并更新参数
    gradients = tape.gradient(weighted_loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return weighted_loss

# 训练循环示例（假设train_dataset是包含(图像,标签,权重)的tf.data.Dataset）
epochs = 50
for epoch in range(epochs):
    print(f"Epoch {epoch+1}/{epochs}")
    total_loss = 0.0
    step_count = 0
    for batch_images, batch_labels, batch_weights in train_dataset:
        loss = train_step(batch_images, batch_labels, batch_weights)
        total_loss += loss.numpy()
        step_count += 1
    print(f"平均损失: {total_loss/step_count:.4f}")

给新手的小建议 #

1. 优先尝试方案1，代码改动最少，最容易验证效果；
2. 注意权重图的缩放：如果权重值过大，可能会导致损失值爆炸，建议先把权重图归一化到[0, 2]或者[0, 5]区间，根据训练情况调整；
3. 可以先拿几个样本手动计算损失，和函数输出对比，确保损失函数的计算逻辑是正确的。

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