Keras中无损实现复杂（多入参）的损失函数

前言

Keras中，直接利用API可以快速的实现一些功能简单的自定义损失函数：

model.compile(loss=YOUR_CUSTOM_LOSS_FUNCTION)

然而任何的简单都是有代价的，通过这个内置方法定义的损失函数有且只能有y_true和y_pred两个入参:

def simple_loss(y_true, y_pred):
    pass

由于Keras的目标是让非编码专业的人士也能接触AI，这样的设计也不是没有道理的，因为这样可以在降低初阶用户使用门槛的同时规避一些乱七八糟的Bug。但是，遵照这样的设计理念(Keras团队原话为：“as designed”），模型中无法直接获取fit_generator()中传入的target(y_true)，导致复杂损失函数在Keras中的实现稍显麻烦（其实也不难）。不过，为了Keras漂亮的进度条，这点麻烦算什么呢?

背景

在复杂的模型设计中，Loss并不能简单的由y_true和y_pred计算出来，这里，我们用近年来著名的Mask-rcnn来帮助理解(细节其实不用多想，只用注意到需求就可以了)：

粗略的来说，Mask-rcnn是由下面三个部分组成的

1. Backbone

前半部分选择Resnet, Xception等任一工作良好的卷积网络用作Feature提取，后半部分利用 Feature Pyramid Network(FPN) 成多尺度的Feature Map

2. Region Proposal Network (RPN)

根据Feature Map来生成感兴趣区域（ROIs）

3. 并行的两个子网络

• ROI Classifier和Bounding Box Regressor 根据RPN提供的ROI中判断并生成最终检测目标的类别（class)与边界框（Bounding box）
• Mask卷积网络 用于生成最后的Mask

这是一个包含多个子模型的复杂模型，#3子模型 的损失函数的在y_true和y_pred之外还需要 #2子模型 输出的ROIs作为入参。 此时，两个入参（y_true, y_pred）的简单损失函数便无法胜任了。

方法

接下来便开始讲解如何 无损的用Keras来构造类似def my_loss(y_true, y_pred, another_input_01, another_input_02, ...)这样的复杂损失函数。这里的无损，指的是相较于苏剑林-科学空间中的方法1，本文所介绍的方法不会损失Keras自带Metrics显示。事实上，本文介绍的方法更像是上述连接中方法的一个完善，但由于本人先受到MatterPort的启发找到解决方法以后再看到的这篇Blog，所以这里便不说是以苏神的想法为参考了。

作为例子，我们首先构造一个简单的网络结构以及一个简单的只有2个参数的自定义loss

from keras import layers as KL
from keras import models as KM

def create_simple_model():
    input_img = KL.Input([64, 64, 3])
    branch1 = KL.Conv2D(64, (3, 3), strides=(4, 4), activation="relu")(input_img)
    branch2 = KL.Conv2D(64, (3, 3), strides=(4, 4), activation="relu")(input_img)
    concat1 = KL.Concatenate()([branch1, branch2])
    deconv1 = KL.Deconv2D(1, (3, 3), strides=(4, 4), activation="relu")(concat1)
    output = KL.Conv2D(1, (1, 1), strides=(1, 1), activation="sigmoid")(deconv1)
    return KM.Model(inputs=input_img, outputs=output)

def my_simple_loss(y_true, y_pred):
    # do what you want here
    return binary_crossentropy(y_true, y_pred)

定义随机生成数据的generator:

import numpy as np

def fake_data_generator(num_samples):
    while (1):
        imgs = np.random.random((num_samples, 64, 64, 3)).astype("float32")
        masks = np.random.random((num_samples, 64, 64, 1)).astype("float32")
        yield imgs, masks

编译模型并开始训练：

train_gen = fake_data_generator(10)
val_gen = fake_data_generator(5)

model = create_simple_model()
model.summary()
model.compile(optimizer="adam", loss=my_simple_loss)
model.fit_generator(
    train_gen,
    epochs=10,
    steps_per_epoch=50,
    validation_data=val_gen,
    validation_steps=5
)

然后Keras经典的实时训练的进度条便出现了：

20190104-1.jpg

接下来便是重头戏了，多个入参的复杂损失函数如何实现呢？

我们首先定义这样一个函数，分别用网络中不同层deconv1, output的输出与y_pred分别求不同的loss然后相加得到最后总的loss(hint：把不同的loss结合起来求一个总的loss是一个很常用的技巧,可以综合不同loss的优点，在Data Science Bowl 2018中，第一名的获得者就是使用了加权的dice loss和bce loss最终得到了令人惊讶的成绩。当然，本文这里的2个loss结合的例子并没有什么道理，只是为了介绍方法，请勿生搬硬套)。

from keras.losses import mean_squared_error, binary_crossentropy

def my_complex_loss_graph(target, deconv1, output):
    mse_deconv1 = mean_squared_error(target, deconv1)
    bce_output = binary_crossentropy(target, output)
    final_loss = mse_deconv1 + bce_output
    return K.mean(final_loss)

有了3个入参的损失函数，我们的模型也必须做相应的更改：

import tensorflow as tf 

def create_complex_model(mode="train"):
    assert mode in ("train", "predict"), "only 'train' and 'predict' mode supported"

    input_img = KL.Input([64, 64, 3])
    branch1 = KL.Conv2D(64, (3, 3), strides=(4, 4), activation="relu")(input_img)
    branch2 = KL.Conv2D(64, (3, 3), strides=(4, 4), activation="relu")(input_img)
    concat1 = KL.Concatenate(name="concat1")([branch1, branch2])
    deconv1 = KL.Deconv2D(1, (3, 3), strides=(4, 4), activation="relu")(concat1)
    output = KL.Conv2D(1, (1, 1), strides=(1, 1), activation="sigmoid")(deconv1)

    if mode == "train":
        #本文最开始提到过，keras generator中yield input, target的target是无法获取的
        #参考github issues:https://github.com/keras-team/keras/issues/11812
        #所以为了取到target,我们必须须把target也当作inputs的一部分传进来即 
        #yield  [input,target], [], 然后再通过KL.Input按顺序获取
        target = KL.Input([64, 64, 1], name="target")
        my_complex_loss = KL.Lambda(
            lambda x: my_complex_loss_graph(*x), name="complex_loss"
        )([target, deconv1, output])
        inputs = [input_img, target]
        outputs = [output, my_complex_loss]
    else:
        #predict阶段，就不用计算loss了所以这里不加入loss层和metric层
        inputs = input_img
        outputs = output

    model = KM.Model(inputs=inputs, outputs=outputs)
    
    #重点
    model._losses = []
    model._per_input_losses = {}
    #通过add_loss来把之前通过KL.Lambda定义的层加入loss，当添加了多个loss层时，optimizer实际优  
    #化的是多个loss的和
    for loss_name in ["complex_loss"]:
        layer = model.get_layer(loss_name)
        if layer.output in model.losses:
            continue
        loss = tf.reduce_mean(layer.output, keepdims=True)
        model.add_loss(loss)
    #其实这里可以添加的不只loss， 有助于监视模型情况的metrics比如f1 score, iou等等也可以通过   
    #model.metrics_tensors.append()来添加

    return model

别被突然增加的代码吓到，其实原理很简单，把loss的计算图通过Lambda转换为layer然后把layer通过add_loss编译进模型，相应的，generator也需修改一下：

def fake_data_generator_2(num_samples):
    while (1):
        imgs = np.random.random((num_samples, 64, 64, 3)).astype("float32")
        masks = np.random.random((num_samples, 64, 64, 1)).astype("float32")
        inputs = [imgs, masks]
        targets = []
        yield inputs, targets

训练：

train_gen = fake_data_generator_2(10)
val_gen = fake_data_generator_2(5)
model = create_complex_model("train")
model.summary()
model.compile(
    optimizer="adam",
    loss=[None] * len(model.outputs)
)
model.fit_generator(
    train_gen,
    epochs=10,
    steps_per_epoch=50,
    validation_data=val_gen,
    validation_steps=5
)

Keras进度条如下：

20190104-2.jpg

最后

当然，如果你如果通过上面代码注释中的方法添加了多个loss和多个metrics的话，你的进度条可能是这样的（这里loss != mask_bce_loss+mask_dice_loss是因为如果把所有loss都显示在进度条上的话会看起来特别凌乱，所以我隐藏了一部分loss，实际上他们还是在工作的）：

20190104-3.jpg

唔，真是...赏心悦目啊! Happy tuning!