TensorFlow2.0的动态图和静态图切换 part 1

TensorFlow2.0的动态图和静态图切换

tf.function介绍

动态图是tf2.0最引人注目的特征，大部分其他改动都是为了适应动态图。它允许将一部分python语法转换为可移植、高性能、语言无关的TensorFlow1.x语法，通过这种方式完成tf1.x静态库和tf2.0动态图的切换。

尽管听起来很美好，但是tf.function还是有一些不为人知的细节值得我们注意，本文就是通过错误驱动的方式来探索tf.funcion的使用技巧。

tf1.x的Session执行

首先我们来回顾一下tf1.x中利用Session来完成图计算的过程：

1. 创建一个tf.Graph对象，并设置为当前范围的默认图
2. 通过TensorFlow API描述这个图的计算过程
3. 提前想好参数共享并定义好各个参数的使用范围
4. 创建并配置好tf.Session
5. 编译静态图，并加载到Session中
6. 初始化参数
7. 使用tf.Session.run法来启动这个计算过程。执行过程将触发一个回溯机制，从选定的节点一直回溯到各层的输入，来决定每个节点的依赖以及计算最终结果。

上面这些过程可以通过以下代码了解：

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))

tf2.0默认的是直接运行(eager execution)模式，和1.x不一样，2.0的运行过程完全和人的想法是一致的。

• 去掉了图的定义
• 去掉了Session的执行
• 去掉了变量的初始化
• 去掉了参数的范围限制
• 去掉了tf.control_dependencies，可以直接运行一个没有依赖关系的操作序列

代码示例如下：

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())

相比于tf1.x，直接执行的缺点就是运行速度慢，因为它依赖python的解释器来运行所有的计算过程，而这是很慢的，且有一些优化仅仅只能用在静态库上。简而言之，开发效率和运行效率不可兼得。

如果我们鱼和熊掌都想要呢？tf2.0 提供了tf.function来兼容静态图和动态图两种模式。

使用tf.function 而不是 tf.Session

tf2.0最大的变化之一就是去除了Session类。这个改变强制开发者使用python函数和一些装饰语句来构造计算图，不能再使用Session来定义了。也就是说我们可以通过tf.function装饰器来加速python函数。

注意：静态图相对动态图的加速是不一定的。对于一些简单的过程，不太值得我们去定义静态图，比如简单的矩阵乘法操作等；但是对于深度神经网络的计算，转换到静态图是可以有比较明显的性能提升的。

从python函数到它的图表示的转换称为AutoGraph。

在tf2.0中，python函数会自动转换到它的图表式，只需要使用@tf.function装饰，这个装饰语句会自动生成python函数的可调用的静态图。

tf.funcion : 简单的解释

在第一次调用使用tf.function装饰的python函数时：

• 函数被执行且被追踪。在这个环境下，eager execution被禁止了，就像tf1.x一样，每个tf的方法都定义了一个tf.Operation，有输入和输出的Tensor对象。
• AutoGraph可以用来检测可以被转换的python的结构体，并转换为相应的图表示(while -> tf.while, for -> tf.while, if -> tf.cond, assert -> tf.assert …)
• 通过函数的追踪+AutoGraph，函数的图表示就有了。为了保持图的计算顺序，tf.control_dependencies会被自动添加到每个执行语句，来定义i+1行和第i行的依赖条件。
• tf.Graph对象被编译
• 一个唯一的ID会给到这个构建好的图，这个ID是根据函数名字和输入参数生成的。这个图会被放到一个map中去，map[id]=graph
• 如果key能够匹配成功，每次函数调用都是在重复使用已经构建好的图

转换到eager execution

为了使用tf.function，我们首先要做的就是重构tf1.x的代码，这个过程是比较容易的，将Session中运行的函数部分提取出来携程python函数：

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

如果我们执行f()，会发生什么？ 什么都不会发生！因为tf2.0默认的eager execution模式会把f当做一个标准的python函数来执行。

从eager到tf.function：需要做的准备

让我们在f前面增加一个@tf.function语句，为了清晰的说明(也为了使用传统打印的方式来debug)，我们还在函数体内增加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()

会发生什么❓

1. @tf.function会将f函数包装成一个tensorflow.python.eager.def_function.Function对象，而这个python函数f，是这个Function对象的.python_function成员变量。
2. 直到这个对象( f( ) )被调用，什么都不会发生
3. 当f()被调用时，图的编译就开始了。此时，只有python代码会被执行且追踪，来获取足够的数据构建图。tf.print不会被执行，因为它会被编译到图中，因此输出是：

PRINT: Tensor(“add:0”, shape=(2, 2), dtype=float32)

4. FAIL：函数第一次执行的时候，会报以下错误，@tf.function编译失败：

ValueError: tf.function-decorated function tried to create variables on non-first call.

原因是：

State (like tf.Variable objects) are only created the first time the function f is called.

第一次调用：

f()

因为是第一次调用，所以f被执行，且图被定义。
** 其他调用 **

f() #again

报错：

ValueError: tf.function-decorated function tried to create variables on non-first call.

这个原因是很好解释的，之所以报错是因为函数定义中包含了一个tf.Variable定义。实际上，在eager模式中，这个对象就是一个普通的python对象，在定义范围之外会被自动回收，然后下次运行会重新定义，因此不会有错误。然而tf.Variable定义了一个持久的对象，如果函数被@tf.function修饰，eager execution被禁止，tf.Variable定义的实际上是图中的一个节点，而这个节点不会被自动回收，且图一旦编译成功，不能再创建变量。

因此，同一个函数在eager 模式下完美运行，但是使用tf.function修饰之后运行报错。这就是我们踩得第一个坑：

要将一个函数从eager 模式转换为它的图表示，我们需要好好想想这个图的编译和运行过程，即便我们当前是在eager模式中。

那么我们怎么避免这个问题呢？ 有三种方法：

1. 定义f() 接受一个输入，这个输入可以是tf.Variable或者其他类型
2. 定义一个全局变量b，在f()中判断b是否是None，如果不是，则不再重复定义
3. 将所有代码包装到一个类中，将b作为一个类成员变量，如果变量已经被初始化，则不再重复定义

丑陋的全局变量

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()

包装成类

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()

包装成类的方式更好：没有全局变量，这个F()类可以随时被实例化且完全不用担心变量b是否会被重复定义。这是我们踩得第二个坑：

当我们尝试将一个函数转换成它对应的加速的图，我们要思考这个图的构建过程。eager模式的python函数和tf.function的函数并不是1:1对应的，由于AutoGraph的存在，我们不需要关注执行顺序，但是我们还是需要关注那些可以创建状态的对象(tf.Variable)。

传参数的方法

@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)

和之前的两种方法一样，传参的方法也可以正常运行，且由于传入的是状态的引用，因此参数在函数内部被修改之后，外部调用时结果也会被修改。

a = tf.Variable(0)

@tf.function
def g(x):
    x.assign_add(1)
    return x

print(g(a))
print(g(a))
print(g(a))

这个会输出

1
2
3

结论

1. tf.function可以作为tf1.x和tf2.0的桥梁，切换静态图和动态图
2. 一般来说，不能直接使用tf.function来修饰一个python函数，我们需要做一些重构来满足条件
3. tf.function修饰的函数不能定义tf.Variable，可以通过封装成类、传参等方法来解决这个问题