keras 如何优雅地实现多个自定义损失函数

keras 如何优雅地实现多个自定义损失函数

之前跑SRGAN模型的时候，由于有多个损失函数（content loss、adversarial loss、perceptual loss）当时采用的是train_on_batch进行训练的。但是发现相比于model.fit，train_on_batch的训练速度会慢很多（虽然我也不知道为啥）而且还无法显示loss变化的进度调，所以就让人很难受QAQ。后来看见知乎上这位作者的文章Tensorflow2.0中复杂损失函数实现，那现在就说一下我是怎么实现多个自定义损失函数的吧。
主要思想是把自定义的损失函数作为一个layer添加到模型中，然后再把自定义的损失作为output输出
以下代码基于tensorflow-gpu==2.5.0完成
首先还是先导包：

import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras import layers
from tensorflow.keras import Model

使用Subclass定义损失函数层（这里的损失函数可输入多个参数，但是我其实只需要两个：y_true，y_predict），我这里一共定义了两个损失函数，分别是Lap_Loss和SSIM_Loss.

# 第一个损失函数
class Lap_Loss(layers.Layer):
    def __init__(self, **kwargs):
        super(Lap_Loss, self).__init__(**kwargs)
     
    def LapPyramid(self, input):
        # 构建拉普拉斯金字塔，进行下采样: 步长为2，先下采样，后上采样，再做差值
        down0 = tf.image.resize(input, [128, 128], preserve_aspect_ratio=True, method="gaussian")
        up0 = tf.image.resize(down0, [256, 256],preserve_aspect_ratio=True, method="gaussian")
        diff0 = Subtract()([input,up0]) # 和原图的差

        down1 = tf.image.resize(down0, [64, 64],preserve_aspect_ratio=True, method="gaussian")
        up1 = tf.image