Tensorflow 2.0上手6: 解剖tf.function的使用

关于Tensorflow 2.0，最令我觉得有意思的功能就是tf.function和AutoGraph了．他们可以把Python风格的代码转为效率更好的Tensorflow计算图．为了更好的了解这两个功能，我在这里分享一下Paolo Galeone对tf.function和AutoGraph介绍的系列文章．

Session execution

了解Tensorflow 1.x的读者都知道Tensorflow一般的工作流程．创建一个计算图tf.graph，然后通过tf.Session对计算图进行计算．以下是一段简单的代码：

g = tf.Graph()
with g.as_default():
    a = tf.constant([[10,10],[11.,1.]])
    x = tf.constant([[1.,0.],[0.,1.]])
    b = tf.Variable(12.)
    y = tf.matmul(a, x) + b
    init_op = tf.global_variables_initializer()
with tf.Session() as sess:
    sess.run(init_op)
    print(sess.run(y))

用Tensorflow 2.0默认的Eager execution实现则很不一样，用户不再需要直接定义计算图或者通过tf.Session来执行代码，也不需要调用tf.global_variables_initializer去初始化变量或者通过tf.control_dependencies去执行计算图中没有包含的节点．

Tensorflow变得像普通的Python代码一样简单．

a = tf.constant([[10,10],[11.,1.]])
x = tf.constant([[1.,0.],[0.,1.]])
b = tf.Variable(12.)
y = tf.matmul(a, x) + b
print(y.numpy())

这样的代码非常易读，但是也带来了执行效率低的问题，因为代码需要依赖Python的解释器来进行计算，无法对数据流以及计算图进行优化．

为了在代码可读性和代码速度之间保持平衡，Tensorflow 2.0引入了tf.function这一概念．

tf.function, not tf.Session

Tensorflow 2.0中一个主要的改变就是移除tf.Session这一概念．这样可以帮助用户更好的组织代码，不用将tf.Session作为一个变量在Python函数中传来传去，我们可以用一个Python装饰符来进行加速，那就是@tf.function．

需要注意的是不是所有的函数都可以通过tf.function进行加速的．有的任务并不值得将函数转化为计算图形式，比如简单的矩阵乘法．然而，对于大量的计算，如对深度神经网络的优化，这一图转换能给性能带来巨大的提升．我们也把这样的图转化叫作tf.AutoGraph．在Tensorflow 2.0中，我们会自动的对被@tf.function装饰的函数进行AutoGraph优化．

懒人介绍tf.function

我们来粗浅的看一下被tf.function装饰的函数第一次执行时都做了什么：

1. 函数被执行并且被跟踪(tracing)．Eager execution处于关闭状态，所有的Tensorflow函数被当做tf.Operation进行图的创建．
2. AutoGraph被唤醒，去检测Python代码可以转为Tensorflow的逻辑，比如while > tf.while, for > tf.while, if > tf.cond, assert > tf.assert．
3. 通过以上两部，我们对函数进行建图，为了保证Python代码中每一行的执行顺序，tf.control_dependencies被自动加入到代码中．保证第i行执行完后我们会执行第i+1行．
4. 返回tf.Graph，根据函数名和输入参数，我们将这个graph存到一个cache中．
5. 对于任何一个该函数的调用，我们会重复利用cache中的计算图进行计算．

我们来看一下Tensorflow 2.0中Eager Execution的代码如何转为tf.function的代码．首先来看一段简单的Tensorflow 2.0代码:

def f():
    a = tf.constant([[10,10],[11.,1.]])
    x = tf.constant([[1.,0.],[0.,1.]])
    b = tf.Variable(12.)
    y = tf.matmul(a, x) + b
    return y
print(f().numpy())
#执行结果
[[22. 22.]
 [23. 13.]]

因为Tensorflow 2.0默认是Eager execution，代码的阅读和执行就和普通的Python代码一样，简单易读．

From eager to tf.function

首先我们简单的加上@tf.function装饰一下，为了方便调试，我们加入一个print和一个tf.print．

@tf.function
def f():
    a = tf.constant([[10,10],[11.,1.]])
    x = tf.constant([[1.,0.],[0.,1.]])
    b = tf.Variable(12.)
    y = tf.matmul(a, x) + b
    print("PRINT: ", y)
    tf.print("TF-PRINT: ", y)
    return y
f()
#执行结果
PRINT:  Tensor("add:0", shape=(2, 2), dtype=float32)
ValueError: tf.function-decorated function tried to create variables on non-first call.

咦，为什么代码报错了？

我以为我发现了一个Bug，因为我理想的执行方式是第一次执行后进行构图并且调用这个图进行计算，之后每次都重复调用这个计算图，然而我们却收到了异常．

但实际上程序的执行过程和我们期待的是一样的，因为tf.function可能会对一段Python函数进行多次执行来构图，在多次执行的过程中，同样的Variable被创建了多次，产生错误．

这其实也是一个很容易混乱的概念，在eager mode下一个Variable是一个Python object，所以会在执行范围外被销毁．但是在tf.function的装饰下，Variable变成了tf.Variable，是在Graph中持续存在的．

把一个在eager mode下正常执行的函数转换到Tensorflow图形式，需要一边思考着计算图一边构建程序．

所以我们在使用tf.function有几种操作需要做：

1. 设计函数f时需要一些输入参数，这个输入参数可以是tf.Variable或者其他任何类型．
2. 设计一个函数从parent scope继承Python variable，在函数中检查Variable是否已经定义过(if b != None)．
3. 将所有的内容写到一个class里，就好像Keras layer一样,所有的Variable都是class的内部参数(self._b)，将class的__call__()通过tf.function装饰．

我们来分析一下上面这些写法．

方案2 vs 方案3

上面的第2种方案和第3种方案类似，但是从面向对象的角度来看，方案3会更美一些．

我们先来看一下方案2的丑陋写法(很不推荐):

b = None
@tf.function
def f():
    a = tf.constant([[10, 10], [11., 1.]])
    x = tf.constant([[1., 0.], [0., 1.]])
    global b
    if b is None:
        b = tf.Variable(12.)
    y = tf.matmul(a, x) + b
    print("PRINT: ", y)
    tf.print("TF-PRINT: ", y)
    return y
f()

方案3的写法:

class F():
    def __init__(self):
        self._b = None
    @tf.function
    def __call__(self):
        a = tf.constant([[10, 10], [11., 1.]])
        x = tf.constant([[1., 0.], [0., 1.]])
        if self._b is None:
            self._b = tf.Variable(12.)
        y = tf.matmul(a, x) + self._b
        print("PRINT: ", y)
        tf.print("TF-PRINT: ", y)
        return y
f = F()
f()

函数逻辑与主过程的分割相对清晰，没有全局变量，class F可以被实例化然后调用，而不用担心我们生成一个全局变量b，让其他函数也能看到．

运用上述两种代码，我们已经解决了对于函数创造自己的Variable的问题．在实际运行过程中，tf.function生成的函数将和eager execution返回相同的结果．

方案1举例

现在来看一下上面讲到的方案1，将变量传递到函数中去．

@tf.function
def f(b):
    a = tf.constant([[10,10],[11.,1.]])
    x = tf.constant([[1.,0.],[0.,1.]])
    y = tf.matmul(a, x) + b
    print("PRINT: ", y)
    tf.print("TF-PRINT: ", y)
    return y
b = tf.Variable(12.)
f(b)

如上一节所述，该函数也可以产生预期的行为。此外，通过传递Variable，我们也可以在函数内部更新Variable的值．比方下面的代码可以更新参数x，产生1,2,3．

@tf.function
def g(x):
    x.assign_add(1)
    return x
a = tf.Variable(0)
print(g(a))
print(g(a))
print(g(a))
#执行结果
tf.Tensor(1, shape=(), dtype=int32)
tf.Tensor(2, shape=(), dtype=int32)
tf.Tensor(3, shape=(), dtype=int32)

Conclusions

关于tf.function今天就先讲到这里，我们大致了解了怎么把一个Python eager mode下执行的函数通过tf.function转化为更快的计算图．

等下一章我们来仔细研究一下当我们把一个tf.Tensor或者Python value传入被tf.function装饰的代码里会发生什么．以及tf.function第一次被执行时，里面的Python代码究竟是如何运行的，每一步都会被转为我们期待的计算图形式吗？

原文: Dissecting tf function part 1
https://pgaleone.eu/tensorflow/tf.function/2019/03/21/dissecting-tf-function-part-1/