Keras中无损实现复杂(多入参)的损失函数

本文基于比较古旧的KERAS=2.1.5版本,运用了最新tf2.0以及tf.keras特性的更好版本请移步我的另一篇文章:

Ziyigogogo:Tensorflow2.0中复杂损失函数实现zhuanlan.zhihu.com图标

前言

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经典的实时训练的进度条便出现了:

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

我们首先定义这样一个函数,分别用网络中不同层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进度条如下:

最后

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

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

编辑于 2019-07-17