tensorflow2.0系列（5）：再说说tf.function

上一篇啰啰嗦嗦说了很多tf2.0中eager execution和autograph的一些特性，但是感觉还是没有说透，反而让人很迷惑，这次再唠唠tf.function到底有啥奇妙之处。

tf.function()

tf1.x中一般的工作流程，就是先创建一个计算图，然后通过tf.Session对图进行计算。例如：

g = tf.Graph()

with g.as_default():

    a = tf.constant(...)

    x = tf.constant(...)

    b = tf.Variable(..)

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

eager execution带来的改变就是，正如大家多次听说的那样，用户不再需要直接定义计算图或者通过tf.Session来执行代码，也不需要调用tf.global_variables_initializer去初始化变量或者通过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())

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

Autograph

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

因此，在被tf.function装饰的函数在执行时，tensorflow编译器会发生下面的一系列操作：
1） 函数被执行并且被跟踪，eager execution处于关闭状态，所有的tf函数都会被当作operation进行图的构建。
2）AutoGraph被唤醒，检测python代码转为tensorflow的逻辑，例如while -> tf.while_loop, for, -> tf.while, if -> tf.cond, assert -> tf.assert。
3）构建图，为了保证代码的执行顺序，tf.control_dependencies被自动加入到代码中，保证第i行执行完后执行第i+1行。
4）返回tf.Graph， 根据函数名和输入参数，将这个graph存入缓存中。
5）对于任何一个对该函数的调用，都会利用缓存中的计算图进行计算。

我们来看下面一个例子：

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:  tf.Tensor(
[[22. 22.]
 [23. 13.]], shape=(2, 2), dtype=float32)
TF-PRINT:  [[22 22]
 [23 13]]

如果加上@tf.function装饰器呢？得到的执行结果却是：

PRINT: Tensor("add:0", shape=(2, 2), dtype=float32)
ValueError: tf.function-decorated function tried to create variables on non-first call.

出错了，而且是在tf.print("TF-PRINT: ", y)这一步，为什么呢？
这是因为tf.function可能会对一段Python函数进行多次执行来构图，在多次执行的过程中，同样的Variable被创建了多次，产生错误．
这其实也是一个很容易混乱的概念，在eager mode下一个Variable是一个Python object，所以会在执行范围外被销毁．但是在tf.function的装饰下，tf.Variable是在Graph中持续存在的．
所以，我们在使用tf.function的时候，应该充分考虑到这个特性：
1）设计函数是需要一些输入参数，这个输入参数可以是tf.Variable或者其它的任何类型；
2）设计一个函数从parent scope继承Python的variable，在函数中检查该variable是否已经被定义过，例如if var != None
3）将所有内容写到一个class里，类似keras layer一样，所有的variable都是class的内部参数（self.b）, 将class的__call__( )通过tf.function装饰。
所以，上面例子的解决方法可以是：

@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


b = tf.Variable(12.)

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 yf = F()

f()

更多阅读

在阅读tensorflow深度学习算法实现的时候，我们经常会看到super.init()这样的代码，这是python OOP类的继承特性。具体可以看：

python类的super和__init__()的关系：

https://www.cnblogs.com/python-nameless/p/6229506.html
https://www.runoob.com/python/python-func-super.html