首页 > 解决方案 > 具有像素加权交叉熵的 U-Net:输入维度错误

问题描述

我一直在使用Zhixuhao 的 U-Net 实现来尝试进行语义二进制分割,并使用 Stackoverflow 答案中的建议对其进行了轻微修改: Keras, binary segmentation, add weight to loss function
to be able to do a pixel-wise weighted binary交叉熵,就像他们在原始 U-Net 论文(见第 5 页)中所做的那样,强制我的 U-Net 学习边界像素。本质上,这个想法是添加一个 lambda 层,该层计算模型本身内的像素加权交叉熵,然后使用仅复制网络输出的“身份损失”。

这是我的输入数据的样子:
输入图像 真实 权重

这是我的代码的样子:

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 的情况下)

标签: pythontensorflowmachine-learningkerasdeep-learning

解决方案


因此,对于偶然发现这个问题的任何人,以下是我实现损失函数的方式:

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)

推荐阅读