tensorflow中同时两个损失函数_Tensorflow2.0中复杂损失函数实现

Tensorflow 2.0自4月初alpha发布以来，引起了广泛关注。其中，谷歌携手

前言

自TF2.0发布以来，我的工作就是把公司之前Tensorflow+Keras(stand-alone)的AI框架转移到Tensorflow2.0中。一方面是希望用到TF2.0的一系列新特性，另一方面，由于之前独立Keras与Tensorflow的工作流程中，遇到过多次版本不匹配带来不易察觉的问题，所以希望能利用上tf.keras来减少环境依赖从而避免这种无谓的坑。

通过数月的持续实践，我惊喜的发现tf.keras并不是直接无脑把独立keras搬进了Tensorflow，谷歌及Keras团队为tf.keras做出的一系列专属优化使得tf.keras无论是在执行性能，模型表现还是易用性上相比独立Keras+Tensroflow的模式都更胜一筹，2个简单的例子：tf.data在构建数据管道(Input Pipeline)的时候，速度以及稳定性都完爆独立Keras中的DataSequence(实际测试中，良好的调优下tf.data.Dataset在model.fit()中数据准备的效率是DataSequence的4倍以上)

tf.distribute在多GPU训练中相对于独立Keras中的multi_gpu_model函数，在显存占用，训练速度，以及最终模型表现中都明显更优。

于是感受到明显技术进步的我决定写文章来记录我几个月以来对TF2.0使用的一些实践。

在上一篇文章中：Ziyigogogo：Keras中无损实现复杂(多入参)的损失函数​zhuanlan.zhihu.com

我介绍过如何用独立Keras库如何实现复杂的多入参损失函数。对比之前的实现，这篇文章将介绍TF2.0中一种更好的方法，使得我们自定义的复杂损失函数可以更容易的在不同的模型架构中重复使用，下面直接上代码：

以下代码基于tf2.0 beta版本实现，安装方法：

pip install tensorflow-gpu==2.0.0-beta1

首先导包：

import tensorflow as tf

from tensorflow.python.keras import backend as K

from tensorflow.python.keras import layers as KL

from tensorflow.python.keras import models as KM

import numpy as np

接下来利用Subclass自定义一个损失函数层：

​

class WbceLoss(KL.Layer):

def __init__(self, **kwargs):

super(WbceLoss, self).__init__(**kwargs)

​

def call(self, inputs, **kwargs):

"""# inputs：Input tensor, or list/tuple of input tensors.如上，父类KL.Layer的call方法明确要求inputs为一个tensor，或者包含多个tensor的列表/元组所以这里不能直接接受多个入参，需要把多个入参封装成列表/元组的形式然后在函数中自行解包，否则会报错。"""

# 解包入参

y_true, y_weight, y_pred = inputs

# 复杂的损失函数

bce_loss = K.binary_crossentropy(y_true, y_pred)

wbce_loss = K.mean(bce_loss * y_weight)

# 重点：把自定义的loss添加进层使其生效，同时加入metric方便在KERAS的进度条上实时追踪

self.add_loss(wbce_loss, inputs=True)

self.add_metric(wbce_loss, aggregation="mean", name="wbce_loss")

return wbce_loss

可以看到，相对于之前使用Lambda把损失函数包装成Layer的写法，我们现在使用了KL.Layer的Subclass写法，看起来似乎代码行数增加了，但是，使用起来却会方便许多：

def my_model():

# input layers

input_img = KL.Input([64, 64, 3], name="img")

input_lbl = KL.Input([64, 64, 1], name="lbl")

input_weight = KL.Input([64, 64, 1], name="weight")

predict = KL.Conv2D(2, [1, 1], padding="same")(input_img)

my_loss = WbceLoss()([input_lbl, input_weight, predict])

​

model = KM.Model(inputs=[input_img, input_lbl, input_weight], outputs=[predict, my_loss])

model.compile(optimizer="adam")

return model

然后我们构建假的数据来实验一下我们的模型是否工作：

def get_fake_dataset():

def map_fn(img, lbl, weight):

inputs = {"img": img, "lbl": lbl, "weight": weight}

targets = {}

return inputs, targets

​

fake_imgs = np.ones([500, 64, 64, 3])

fake_lbls = np.ones([500, 64, 64, 1])

fake_weights = np.ones([500, 64, 64, 1])

fake_dataset = tf.data.Dataset.from_tensor_slices(

(fake_imgs, fake_lbls, fake_weights)

).map(map_fn).batch(10)

return fake_dataset

​

​

model = my_model()

my_dataset = get_fake_dataset()

model.fit(my_dataset)

然后就是熟悉的keras进度条了：

50/50 [==============================] - 1s 24ms/step - loss: 7.8311 - wbce_loss: 7.8311

​

# 可以根据需求，把多个自定义loss层加入模型:

554/554 [==============================] - 402s 725ms/step - loss: 0.3199 - wbce_loss: 0.0681 - dice_loss: 0.2518

# 其中loss的数值就代表多个自定义loss的和: 0.3199 = 0.0681 + 0.2518

可以看到，相比

结束语

本文介绍了TF2.0中一种比之前更加便捷的复杂损失函数的写法，同时代码中刻意引出了利用tf.data来构建了简易的数据管道的列子。关于tf2.0中tf.data的详细用法以及最佳实践将会在下一篇文章中详细介绍。