tf eager 模式 --- tf.function 与 autograph

tf eager 模式 — tf.function 与 autograph

declarative vs imperative programing

目前主流的深度学习框架的编程范式可以归为两类。
第一类是 declarative programing。用户首先定义好一个计算图，然后再执行这个计算图。在执行之前可能会做一些图优化。这其实有点像编译型语言的编码、编译优化、执行的流程。tensorflow 2.x 之前的编程范式属于这一类。第二类是 imperative programing。用户不需要提前定义好整个计算，每走一步（执行一个 op),就立马能够得到结果。tf eager 模式属于这一类，tensorflow 2.x 默认开启 eager 模式。

这两种范式各有优劣。declarative programing 提前定义好了整个计算，在执行整个计算之前可以做一些优化，会有一些性能的收益；定义好了计算逻辑，也方便模型部署。然而，这对模型开发和调试却是不友好的。打个比方，假设你在 session run 出错了，报错信息纵然有 op name, 如果你对 tensorflow 底层的计算图不甚熟悉的话，你也不太能够定位到是原始的构图逻辑的哪一行导致了错误。
imperative programing 让用户能够逐步开发、调试模型，能够立马得到反馈，无疑提升了模型开发、调试的效率。但是，受限于逐步执行，imperative programing 无法对整个计算做优化，也无法做 op 的并行调度。

有没有办法能够综合两种编程范式的优势呢？that is where tf.function comes in.

tf.function

tf.function 是个decorator, 用来修饰 function, 对被这个decorator修饰过的函数的调用会生成一个特殊的 op并执行这个 op, 这个 op 的执行会调度执行函数内部定义的计算子图。
简单地说，tf.function 提供了一种把一些计算打包成一个计算图子图在图模式执行的能力。主要用来解决 eager 模式的下性能不足问题。

另外， tf.function的一个附带作用是在图里引入了一定的封装能力，比如， inference 的时候可以把 一个 tf.function 替换成更高效的实现。

demo

以下是一段 tf.function 的 demo 代码

import tensorflow as tf
import numpy as np
#tf.compat.v1.disable_eager_execution()

class F1():
    def __init__(self):
        pass
    # tf.funtion is a decorator    
    @tf.function
    def __call__(self, a, x, v):
        y = tf.matmul(a, x) + v
        return y
        
class F2():
    def __init__(self):
        pass
    @tf.function
    def __call__(self, v):
        a = tf.constant([[10, 10], [11., 1.]])
        x = tf.constant([[1., 0.], [0., 1.]])
        f = F1()
        y = f(a, x, v)
        print("PRINT: ", y)
        # tf.print("TF-PRINT: ", y)
        return y

v = tf.Variable(12.)
f = F2()
t = f(v)
print("PRINT: ", t)
"""
# 在图模式下，查看计算图
import tensorflow._api.v2.compat.v1 as tf
g=tf.get_default_graph()._as_graph_def()
print(g)
"""

这里有一个有意思的事情是，当一个 tf.function 包装的函数调用另一个 tf.funtion 包装的函数时，F2 和 F1 分别运行在 eager 模式和图模式。感兴趣的同学可以加个断点 trace 一下。也就是说，F2 返回的是值， F1 返回的是一个符号。

这里可以这样理解，因为F2 的作用就是把一个子图封装成 一个 op, 然后再一次性地执行这个子图。进入 F2 函数体里面，其实已经是图模式了，每执行一个 tensorflow 的 op, 其实是往子图里添加一个节点而已，并没有实际调用这个 op 的 对应的kernel。

autograph

autograph 和 tf.funtion 一起使用，其作用是把pythonic 的代码转成图模式构图的代码，尽量尽量抹平图模式和eager模式之间的barrier，使同一份代码可以在eager模式和图模式之间顺滑切换。

autograph 主要做了两件事情：

1、控制流转换

autograph 会把 python 的控制流转成图模式计算图里的控制流，即把 if 。。else, while, for之类的控制结构转成对等的tf.cond、tf.while 之类的。
这个转换的前提是条件语句的表达式里面使用到了 tensor.

2、添加控制依赖

autograph 会在节点之间加上必要的控制依赖。比如，对同一个 variable 有读写的op，在 eager 模式，它们是按照 python 语句的顺序执行的；在图模式下，因为有 op 并行调度，为了保证variable 读写顺序和 eager 模式一致，需要在对同一个 variable 读写的 op 之间加上 control dependence。

这两件事情，在python语言内不太可能实现，比如if 。。else, while, for之类，这些是语言内在的结构，并不是运算符，无法重载。需要在语言外部对 tf.function 修饰函数的源代码做修改。
看过 domain specific language 这本书的同学可能有点熟悉，这其实是个语言转换的问题。一般的套路是经过词法分析、语法分析，得到抽象语法树（ast）, 对 ast 做修改，然后把 ast 转成目标语言。tf 实现 autograph 大致也是这个思路。它通过python的反射机制，拿到 function 的源代码，然后使用第三方库解析得到 ast, 最后 ast 生成新的 function 代码。