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.resize(down1, [128, 128],preserve_aspect_ratio=True, method="gaussian")
        diff1 = Subtract()([down0, up1]) # 和原图/2的差

        down2 = tf.image.resize(down1, [32, 32],preserve_aspect_ratio=True, method="gaussian")
        up2 = tf.image.resize(down2, [64, 64],preserve_aspect_ratio=True, method="gaussian")
        diff2 = Subtract()([down1, up2])  # 和原图/4的差

        return diff0,diff1,diff2
​
    def adaptive_loss(self,t,p,scale):
        mse_loss = tf.abs(t-p)
        panduan = mse_loss > scale
        large = mse_loss+scale*10
        result = tf.where(panduan, large, mse_loss)
        out = tf.reduce_mean(tf.square(result))
        return out

    def call(self, inputs, **kwargs):
        """
        # inputs：Input tensor, or list/tuple of input tensors.
        如上，父类layers的call方法明确要求inputs为一个tensor，或者包含多个tensor的列表/元组
        所以这里不能直接接受多个入参，需要把多个入参封装成列表/元组的形式然后在函数中自行解包，否则会报错。
        """
        # 解包入参
        y_true, y_pred = inputs
        # 复杂的损失函数
        t_diff0, t_diff1, t_diff2 = self.LapPyramid(y_true)
        p_diff0, p_diff1, p_diff2 = self.LapPyramid(y_pred)
        
        diff2_loss = 1 - tf.image.ssim(t_diff2, p_diff2,max_val=2)
        diff1_loss = self.adaptive_loss(t_diff1,p_diff1,0.05)
        diff0_loss = self.adaptive_loss(t_diff1, p_diff1, 0.02)
        Lap_loss = diff2_loss + 0.1*diff1_loss + 0.1*diff0_loss
        # 重点：把自定义的loss添加进层使其生效，同时加入metric方便在KERAS的进度条上实时追踪
        self.add_loss(Lap_loss , inputs=True)
        self.add_metric(Lap_loss , aggregation="mean", name="Lap_loss")
        return Lap_loss 

# 第二个损失函数
class SSIM_image_Loss(layers.Layer):
    def __init__(self, **kwargs):
        super(SSIM_image_Loss, self).__init__(**kwargs)

    def call(self, inputs, **kwargs):
        y_true, y_pred = inputs
        ssim_loss = 1 - tf.image.ssim(y_true, y_pred,max_val=2)
        img_ssim = ssim(y_true,y_pred)
        self.add_loss(ssim_loss, inputs=True)
        self.add_metric(img_ssim, aggregation="mean", name="ssim")
        return ssim_loss

在定义完需要的损失函数后就可以构建模型了。

def comb_model():
		# 这里需要给Input定义名字，后面构建数据输入需要用到
        lr = Input(shape=(None, None, 3),name="lr_img")
        hr_image = Input(shape=(None, None, 3),name="hr_img")
        #one_gpu_model是之前写好的超分模型
        sr_image = one_gpu_model(lr) 
        loss_1 = PyramidLoss()([hr_image,sr_image])
        loss_2 = SSIM_image_Loss()([hr_image,sr_image])
        model = Model(inputs=[lr, hr_image], outputs=[sr_image, loss_1, loss_2])
        model.compile(optimizer=Adam(lr=params['lr_init']))
        return model

由于我已经将自定义的两个loss函数以layer的方式加入到model中，因此在complie的时候我只对optimizer做了指定。可以看到，现在model的input=[lr, hr_image]，那么在进行训练时就需要用到：model.fit(x=input_data, y=None),因此需要对输入的数据进行更改。下面对更改前的数据生成代码和更改后的代码进行对比。（如果对tf的数据读入不太清楚的话，可以参考我之前的文章tensorflow数据读入的基础步骤）

# 更改前的数据生成代码
def get_dataset(lr_path, hr_path, ext):
    lr_sorted_paths = get_sorted_image_path(lr_path, ext)
    hr_sorted_paths = get_sorted_image_path(hr_path, ext)

    # 打包hr和lr的所有地址 并进行shuffle打乱
    lr_hr_sorted_paths = list(zip(lr_sorted_paths[:], hr_sorted_paths[:]))
    random.shuffle(lr_hr_sorted_paths)
    lr_sorted_paths, hr_sorted_paths = zip(*lr_hr_sorted_paths)
    # 将hr和lr组合成元组形式
    ds = tf.data.Dataset.from_tensor_slices((list(lr_sorted_paths), list(hr_sorted_paths)))

    def load_and_preprocess_lr_hr_images(lr_path, hr_path, ext=ext):
        return load_and_preprocess_image(lr_path, ext), load_and_preprocess_image(hr_path, ext)
        
    lr_hr_ds = ds.map(load_and_preprocess_lr_hr_images, num_parallel_calls=4)
    return lr_hr_ds, len(lr_sorted_paths)
   # 后续的.shuffle()，.batch()等操作在此省略了

此时，如果使用model.fit(lr_hr_ds)进行训练，keras的内部机制会直接让x=lr,y=hr，这样就无法满足input=[lr, hr_image]的要求。那么下面我们来看更改后的代码：

def load_and_preprocess_lr_hr_images(lr_path, hr_path, ext=ext):
        lr = load_and_preprocess_image(lr_path, ext)
        hr = load_and_preprocess_image(hr_path, ext)
        inputs = {"lr_img":lr,"hr_img":hr}
        targets = {}
        return inputs, targets

    lr_hr_ds = ds.map(load_and_preprocess_lr_hr_images, num_parallel_calls=4)

代码变动的地方只有load_and_preprocess_lr_hr_images这个函数。这里采用字典的形式，将lr和hr放入同一个字典中，需要注意字典中key的命名是要和model中Input的命名一致的。lr = Input(shape=(None, None, 3),name="lr_img") hr_image = Input(shape=(None, None, 3),name="hr_img")。这时运行model.fit(x=lr_hr_ds, y=None)就能够满足input的多输入要求了。
下面就是熟悉的keras进度条啦，因为我定义了两个损失函数，这里的loss其实是Lap_Loss和SSIM_Loss之和。

虽然使用Layer.Subclass的写法看起来代码量会比较多，但是代码的完整度会高很多，后续用起来也更方便。

参考文章
Tensorflow2.0中复杂损失函数实现