具有像素加权交叉熵的 U-Net：输入维度错误

Question

我一直在使用Zhixuhao's implementation of U-Net 尝试进行语义二进制分割，并使用此 Stackoverflow 答案中的建议对其进行了略微修改： Keras, binary segmentation, add weight to loss function
能够像在原始 U-Net 论文 (see page 5) 中那样进行像素加权二进制交叉熵，以强制我的 U-Net 学习边界像素。本质上，这个想法是在模型本身中添加一个计算像素加权交叉熵的 lambda 层，然后使用仅复制网络输出的“身份损失”。

这是我的输入数据的样子：
input image groundtruth weights

这是我的代码的样子：

def unet(pretrained_weights = None,input_size = (256,256,1)):

    inputs = Input(input_size)
    # [... Unet architecture from Zhixuhao's model.py file...]
    conv10 = Conv2D(1, 1, activation = 'sigmoid', name='true_output')(conv9)

    mask_weights = Input(input_size)
    true_masks = Input(input_size)
    loss1 = Lambda(weighted_binary_loss, output_shape=input_size, name='loss_output')([conv10, mask_weights, true_masks])

    model = Model(inputs = [inputs, mask_weights, true_masks], outputs = loss1)
    model.compile(optimizer = Adam(lr = 1e-4), loss =identity_loss)

并添加了这两个功能：

def weighted_binary_loss(X):
    y_pred, weights, y_true = X
    loss = keras.losses.binary_crossentropy(y_pred, y_true)
    loss = multiply([loss, weights])
    return loss

def identity_loss(y_true, y_pred):
    return y_pred

最后是我的 main.py 的相关部分：

input_size = (256,256,1)
target_size = (256,256)
myGene = trainGenerator(5,'data/moma/train','img','seg','wei',data_gen_args,save_to_dir=None,target_size=target_size)
model = unet(input_size=input_size)
model_checkpoint = ModelCheckpoint('unet_moma_weights.hdf5',monitor='loss',verbose=1, save_best_only=True)
model.fit_generator(myGene,steps_per_epoch=300,epochs=5,callbacks=[model_checkpoint])

现在这段代码运行良好，我可以训练我的 U-Net 并且它确实学习了边框像素，但只有如果我将输入图像的大小调整为 256*256。如果我改为在 main.py 中使用 input_size=(256,32,1) 和 target_size=(256,32) ，这是我的数据的相关维度并且允许我使用更大的批量大小，我会收到以下错误：

ValueError：操作数无法与形状一起广播（256， 32, 1) (256, 32)

对于loss = multiply([loss, weights]) 行。事实上，权重有一个额外的单维。我不明白为什么在使用 256*256 输入时不会引发错误，但我尝试使用 k.expand_dims() 或 Reshape() 使两个输入的尺寸相同，但是代码没有发出错误并且损失收敛，当我在额外输入上测试我的网络时，我得到空白输出（即完全灰色或白色或黑色的图像，或与我的输入无关的东西）。

因此，对于以下问题，有很多文字：为什么 multiply() 在 256*32 而不是 256*256 的情况下会发出错误，为什么在输入上创建/删除尺寸没有帮助？

谢谢！

ps：为了让网络在训练后输出实际预测而不是逐像素损失，我使用以下代码删除了损失层和两个额外的输入层：

new_model = Model(inputs=model.inputs,outputs=model.get_layer("true_output").output)
new_model.compile(optimizer = Adam(lr = 1e-4), loss = 'binary_crossentropy')
new_model.set_weights(model.get_weights())

这很好用（至少在 256*256 的情况下）

Answer 1

所以对于任何偶然发现这个问题的人来说，我是这样实现损失函数的：

def pixelwise_weighted_binary_crossentropy(y_true, y_pred):
    '''
    Pixel-wise weighted binary cross-entropy loss.
    The code is adapted from the Keras TF backend.
    (see their github)
    
    Parameters
    ----------
    y_true : Tensor
        Stack of groundtruth segmentation masks + weight maps.
    y_pred : Tensor
        Predicted segmentation masks.

    Returns
    -------
    Tensor
        Pixel-wise weight binary cross-entropy between inputs.

    '''
    
    try:
        # The weights are passed as part of the y_true tensor:
        [seg, weight] = tf.unstack(y_true, 2, axis=-1)

        seg = tf.expand_dims(seg, -1)
        weight = tf.expand_dims(weight, -1)
    except:
        pass

    epsilon = tf.convert_to_tensor(K.epsilon(), y_pred.dtype.base_dtype)
    y_pred = tf.clip_by_value(y_pred, epsilon, 1. - epsilon)
    y_pred = tf.math.log(y_pred / (1 - y_pred))

    zeros = array_ops.zeros_like(y_pred, dtype=y_pred.dtype)
    cond = (y_pred >= zeros)
    relu_logits = math_ops.select(cond, y_pred, zeros)
    neg_abs_logits = math_ops.select(cond, -y_pred, y_pred)
    entropy = math_ops.add(relu_logits - y_pred * seg, math_ops.log1p(math_ops.exp(neg_abs_logits)), name=None)
    
    # This is essentially the only part that is different from the Keras code:
    return K.mean(math_ops.multiply(weight, entropy), axis=-1)