Tensorflow 2.0自4月初alpha发布以来，引起了广泛关注。其中，谷歌携手@fchollet(Keras作者)及其团队对Keras库做出了大量Tensorflow专属的优化以及改动。再联想到独立(Stand alone)的Keras库最近一次更新2.2.4已经是大半年（2018年10月）以前的事情了，不禁八卦Keras团队的工作重心是不是从独立Keras转向了tf.keras来对抗Pytorch的竞争了呢？
前言
自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可以看到，相比上一篇文章中，在model构建之后手动向model._losses的私有属性中添加loss这种偏hack的方法，现在的实现更加优雅方便，可以说和即插即用相差无几了。但是正如上一篇文章中提到的按照Keras作者的设计：我们依然无法获取fit()中传入的target。所以需要把target和input一起传进来。所以，真正像pytorch中那样完全没有额外步骤的损失函数在这个限制开放以前，是无法在keras中实现的。
结束语
本文介绍了TF2.0中一种比之前更加便捷的复杂损失函数的写法，同时代码中刻意引出了利用tf.data来构建了简易的数据管道的列子。关于tf2.0中tf.data的详细用法以及最佳实践将会在下一篇文章中详细介绍。