最近在学习 tensorflow 2.0 (以下tf2)并尝试改写一些1.x时代的代码。在从1.x转向2.0过程中不可避免地遇到了一些坑，其中大多数来源于建图方式的改变，有些得到了解决有些则还在摸索。2.x版本目前还比较新，在此记录下遇到的问题，希望可以为以后的使用提供参考。在阅读本文前强烈建议阅读 tf2 官方指南，尤其是Keras部分。
1. 概述
在2.0中最大的改动就是 session 和显式定义 graph 的删除，并同时力推通过 keras 的自定义对象来实现任意模型，因此 keras 在 tf2 中有举足轻重的地位。而伴随其的另一个重大改动是默认开启 eager execution 环境, 在该环境下，所有 ops 操作会立即执行并计算出具体的值，而不再像以前那样要先构建一个完整的计算图然后再运行。显然，tf2 强调了代码的灵活性，从而使 tf 具有 define by run 的特点, 来和 Pytorch 竞争研究实验领域的地位。但在工业界，运行速度依然很重要，而图结构大大优化了计算速度，仍然是不可或缺的刚需。那在 tf2 中，该如何实现图呢，目前据我所知有两种方式：1）tf.function 装饰器；2）tf.keras.Model 类。具体实现请参考官方指南。可以看出，tf2 意图自然而然地在 eager execution 的基础上添加图的构建，为此加入了AutoGraph的功能以将python语法编写的函数自动转化为图，但是这种隐式的图构建手法对函数和类的设计提出了更高的要求，以便保证默认的 eager execution 部分与图的部分不会产生冲突，对于AutoGraph的各种限制的介绍可以参考这里。
2. @tf.function装饰器
Tf2 中要求用户将图写作一个函数，原本为定义图中的待输入变量而使用的 tf.placeholder 不复存在，而改为利用函数的参数来输入变量值，并通过函数的返回来获得图的计算结果。我们都知道，图结构一旦建立就是固定且封闭的，以前通过 session 可以运行图并获取任意图中变量的值，但现在我们只能获取函数的返回值，而其余的值被函数给隔离开了。那问题是这种隔离是否能被打破呢，让我们来尝试一下。
2.1. 利用Python特性初始化
首先可以想到的是python的函数里的变量并不是完完全全是局部变量，假如某外部变量的同名变量出现在该函数里，同时函数不对其做任何赋值操作，那在调用函数时该变量可以延续其在函数外部的身份，比如：1
2
3
4
5
6
7
8
9def func(x):
y = a + x
return y
a = 1
func(1)
2
a = 2
func(1)
3
假如将 func 装饰为图，结果会怎么样？1
2
3
4
5
6
7@tf.function
def func(x):
y = a * x
return y
a = 1
func(1)
<tf.Tensor: id=1265, shape=(), dtype=int32, numpy=1>
很好，没报错，结果也是2，这看上去似乎很ok，让我们给a换个值：1
2
3a = 2
func(1)
<tf.Tensor: id=1265, shape=(), dtype=int32, numpy=1>
不ok了，输出没有变。可见初次调用图函数时外部变量会发挥作用，但图一旦建立，内外变量的联系便切断了。
当我们换一个输入：1
2
3a = 2
func(2)
<tf.Tensor: id=1271, shape=(), dtype=int32, numpy=4>
a又起作用了，那是因为新的图被建立了。
tf2 中新建一张图的标准有两个：
- 输入中的数值型变量发生改变
- 输入中的张量形状发生改变
我们可以验证一下：1
2
3
4
5
6
7
8
9
10
11
12
13a = 1
x = tf.convert_to_tensor([1,1])
func(x)
<tf.Tensor: id=1298, shape=(2,), dtype=int32, numpy=array([1, 1])>
a = -1
x = tf.convert_to_tensor([2,2])
func(x)
<tf.Tensor: id=1300, shape=(2,), dtype=int32, numpy=array([2, 2])> #输入改变了，但形状没变，因此没有新建图，
#改变的a没有被传入
a = -1
x = tf.convert_to_tensor([2,2,2])
func(x)
<tf.Tensor: id=1308, shape=(3,), dtype=int32, numpy=array([-2, -2, -2])> #输入的形状改变，新建图，a再次发挥作用
由此我们可以知道，在调用图函数前，可以对函数内的变量进行初始化，而一旦建立了图，这些变量就不应该再在外部被改变，一是因为改变了对图也没有作用，二是在下次新建图时会产生不可预知的影响。
另外，通过声明global变量或者将图函数作为某个类的成员函数，可以做到对全局变量或类成员变量赋值，但这同样也是在新建图时一次性有效的。
2.2. 利用tf.summary记录内部变量
1.x中 summary 可以在 session 之外独立获取图内变量值并记录为log，但在2.0我目前还没发现可以这样做的方式。唯一有效的做法是将summary写在图函数中，用法可以参考官方文档，但是我发现官方文档有一处误导：1
2
3
4
5
6
7
8
9
10
11writer = tf.summary.create_file_writer("/tmp/mylogs")
def my_func(step):
# other model code would go here
with writer.as_default():
tf.summary.scalar("my_metric", 0.5, step=step)
for step in range(100):
my_func(step) # Wrong!
writer.flush()
就如我们刚才所验证的，图函数一旦传入不同的数值型参数就会新建图，在这里显然my_func(step)在不断地新建图，这样会使其远远慢于直接eagerly地使用函数。正确的做法应该是把 step 转化为标量张量：1
2
3for step in range(100):
my_func(tf.convert_to_tensor(step,tf.int64))
writer.flush()
注意一定要指定类型为 tf.int64，因为默认转化类型 tf.int32 会使 SummaryWriter 报错。
3. Keras 类
现在一个最大的感受就是 keras 被糅进了 tf 的日常使用当中。当 keras 拥有了可以高度自定义的类，与其说它是 tf 的高层组件，不如说它提供了一个面向对象的建模手段。总得来说，keras提供了层(layer)级的抽象和模型(model)级的抽象，使用者可以将内建的层或者自定义层在模型中任意组合。
事实上，编写自定义层的过程和编写图函数的过程并没有太大区别，而假如对训练过程没有什么特殊需求，完全可以通过 compile 和 fit 轻松实现训练而不必自己写训练过程， keras 的模型在 compile 时是默认建图的，同时 keras 的模型还支持轻松地保存/载入参数和保存/载入模型。但是图函数风格的模型会更具可控性、更加易于调试。在图函数里你可以直观地编写代码而不必关心keras类那些层层包裹的复杂接口。自建图函数 or 使用 keras 类？这个问题官方在 when to use the functional api 有过简短的讨论。
我个人参考教程对 keras 类的自定义做了些探索，其自定义要素体现在这几方面：
- model派生类的自定义
- layer派生类的自定义
- 损失函数的自定义
- 训练过程的设置
越复杂越深度的自定义越容易遇到不可预知的问题，这就要求对这些类有深入的了解。关于自定义layer和model的基础可以参考这部分教程，自定义损失函数则在这部分有叙述。
3.3. 自定义复杂损失函数
从上述教程中可知，自定义损失函数有两种方法：
- 定义一个以真实值和预测值为参数的函数
func(true, pred) - 定义一个 tf.keras.losses.Loss 的派生类，可以在实例化时传递真实值和预测值以外的参数
在用第一种方法时，我遇到了一个问题，那就是当我在 model.fit() 中设置 sample_weight 时，会报如下错误
InvalidArgumentError: Can not squeeze dim[0], expected a dimension of 1, got 1024
其中1024是我的 batch size。其实原本我就很疑惑，假如根据教程我的损失函数返回的是一个标量，那这个 sample_weight 是如何神奇地在样本粒度起作用的。我曾尝试使损失函数接收 sample_weight 变为 func(y_true, y_pred, sample_weight) ，依然报错。结果当我令原本的损失函数输出一个 shape 为 (batch_size,) 的张量时，一切就正常了。我们可以合理推测，sample_weight 并不会传递给损失函数，而是和损失函数得到的结果相乘然后取和作为最终的损失。而当查看 tf.keras.losses.Loss 的文档时发现，它的 __call__() 函数有三个参数 y_true, y_pred 和 sample_weight， 而 call() 则只有 y_true, y_pred，因此可以推测我们定义的损失函数可能是被当作一个 tf.keras.losses.Loss 实例的 call() 函数来使用，在调用该实例时，__call__() 会接收 sample_weight 并调用 call()，将 sample_weight 与 call() 返回的结果相乘并求和然后返回最终结果。
可能有空时我会去看看源码来验证具体的机制。目前的结论就是：当设置有 sample_weight 时，损失函数应该输出一个 shape 为 (batch_size,) 的张量。使用第二种方法时道理也是一样的，只是损失函数变成了 call() 函数。
假如想在训练过程中，获取内部变量作为损失的部分，可以在自定义 layer 时调用 self.add_loss(loss_value)。假如是用 model 的 fit() 函数来训练，则添加的损失会被自动加进最终损失中。假如是自己写的训练过程，则可以用 model.losses 来获取添加的损失，注意 model.losses 是一个列表，单次计算中每次调用 add_loss(loss_value) 都会把 loss_value 添加到列表中，每次计算列表都会被清空。
3.2. trainable_weights 的消失
先来看一个自定义 layer 的例子：1
2
3
4
5
6
7
8
9
10
11
12
13
14class layer(tf.keras.layers.Layer):
def __init__(self, n_features, order):
super(layer, self).__init__()
self.w = [None] * n_features
for i in range(1, order + 1):
rnd_weights = tf.random.uniform((n_features,1), minval = -0.01, maxval = 0.01)
self.w[i - 1] = tf.Variable(0.1, trainable=True, name='coefficient_' + str(i))
self.b = tf.Variable(0.1, trainable=True, name='bias')
def call(self, inputs):
outputs = 0 + self.b
for i in range(1, order+1):
outputs += tf.matmul(tf.pow(inputs, i),self.w[i-1])
return outputs
在这里我根据输入 order 动态构建了了 w 列表内的可训练参数。可是当我最终通过 layer.weights 或者 layer.trainabel_weights 查看参数时却只有 bias, 也就是 b。可见文中 for 循环定义可训练参数会影响模型拾取参数的过程。这种情况到底应该怎么处理最好我也拿不准，希望有高人指点。我的办法是改用 add_weight 接口：1
2
3
4
5
6
7
8class layer(tf.keras.layers.Layer):
def __init__(self, n_features, order):
super(layer, self).__init__()
self.w = [None] * n_features
for i in range(1, order + 1):
rnd_weights = tf.random_uniform_initializer(minval = -0.01, maxval = 0.01)
self.w[i - 1] = self.add_weight(initializer=rnd_weights, trainable=True, name='coefficient_' + str(i), shape=(n_features,1))
self.b = tf.Variable(0.1, trainable=True, name='bias')
注意 add_weight 的 initializer 只接受不带参的 callable，假如想用固定值，lambda:value 是一个简单的方法。
3.4. 利用tf.data.Dataset
关于 tf.data.Dataset 的详细教程可以在这里看到。Dataset 旨在用来接受和转换不同类型的数据源，并提供一系列简单的数据预处理处理手段，并最终作为 model.fit() 的输入。但值得注意的是它和 model.fit() 的参数存在两个冲突：
- batch_size
- sample_weight
也就是说一旦使用了 dataset ， 就不能在 fit 时设定这两个参数，而必须得在 dataset 上动手。
这两个功能在 dataset 上的实现分别为：
train_datase = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)- 将sample_weight作为新建 dataset 时传入的 tuple 的第三个元素，即：
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train, sample_weight))
(可参考这一段)
4.待解决问题
4.1. 图函数对SparseTensor的兼容问题
在使用过程中我发现图函数对SparseTensor并不完全兼容，而且处理SparseTensor的速度上和1代相比有很大差距，这可能是 autograph 还尚不完全成熟，因为对于SparseTensor来说即使是同样的 shape，其 index 和 value 的大小也会不同。考虑过将 SparseTensor 转化为 dataset，但 dataset 的训练必须用 model.fit() ，过程中会出现如下问题。
4.2. 使用 model.fit() 训练时的权重消失问题
3.2里虽然通过 add_weight() 使权重得以自动加入模型的 trainable_weights ，但是用 model.fit() 训练模型时依然没法训练那些权重 ，因此不得不自写训练 loop ，而4.1的问题只能靠特定情况下转为 eager execution 来解决。