Keras：自定义损失取决于另一个输出

Question

我正在尝试实施this论文。

本文中描述的损失函数取决于图像的输出和图像产生的噪声版本。基本模型在图像和噪声图像上都运行，并且损失函数使用噪声图像输出，但是没有直接与地面真实情况进行比较：

def build_model(base_model_name, n_label, noise_std, alpha):
    image_size = get_image_size(base_model_name)
    image_shape = (image_size, image_size, 3)
    image = Input(shape=image_shape)

    noise = layers.GaussianNoise(stddev = noise_std)
    noisy_image = noise(image)

    base_model = get_base_model(base_model_name, n_labels)
    default_output = base_model(image)
    noisy_output = base_model(noisy_image)

    model = Model(inputs=image, outputs=default_output)

    def stability_loss(noisy_output):
        def loss(y_true,y_pred):
            cce = losses.CategoricalCrossentropy()(y_true, y_pred)
            stab = losses.KLDivergence()(y_pred, noisy_output)
            return cce + (alpha * stab)
        return loss

    model.compile(optimizer='adam', loss=stability_loss(noisy_output), metrics=['accuracy'])
    utils.plot_model(model, 'stability_model.jpg', show_layer_names=True, show_shapes=True)
    return model

使用Tensorflow 2.3.0，我能够毫无问题地训练该模型。但是当我使用plot_model打印模型时，我得到以下信息：

因此，噪声层和噪声输出未在模型描述中显示。我想知道的是，在训练过程中会正确生成嘈杂的输出吗？还是我做错了什么？

编辑1：

所以我想我可能已经明白了，我以前的代码似乎已经过时，需要在Tensorflow 2.3.0中取消急切的执行。

相反，我尝试使用Layer.add_loss()方法创建一个增加稳定性调节损失的层：

class StabilityLayer(layers.Layer):
    def __init__(self, alpha, **kwargs):
        super(StabilityLayer, self).__init__(**kwargs)
        self.alpha = alpha

    def get_config(self):
        config = super().get_config().copy()
        config.update({
            'alpha' : self.alpha
            })
        return config

    def call(self, x):
        default_output=x[0]
        noisy_output=x[1]
        kld = self.alpha * losses.KLDivergence()(default_output, noisy_output)
        self.add_loss(kld)
        return default_output

此层已添加到模型：

def build_model(base_model_name, n_label, noise_std, alpha):
    image_size = get_image_size(base_model_name)
    image_shape = (image_size, image_size, 3)
    image = Input(shape=image_shape)

    noise = layers.GaussianNoise(stddev = noise_std)
    noisy_image = noise(image)

    base_model = get_base_model(base_model_name, n_labels)
    default_output = base_model(image)
    noisy_output = base_model(noisy_image)
    output = StabilityLayer(alpha)([default_output,noisy_output])

    model = Model(inputs=image, outputs=output, name='stability_model')
    model.compile(optimizer='adam', loss="categorical_crossentropy", metrics=['accuracy'])
    utils.plot_model(model, 'stability_model.jpg', show_layer_names=True, show_shapes=True)
    return model

这似乎可以创建正确的模型

而且，如果我理解正确，则损失应为： CCE_LOSS + alpha x KLD 如本文中所述，因为Layer.add_loss(...)的损失应自动添加到model.compile(...)