TensorFlow2.X——tf.function && Autograph介绍以及使用

tf.function && Autograph介绍以及使用

要获得最佳性能并使模型可在任何地方部署，可以使用tf.function从程序中构建图。 因为有AutoGraph，可以使用tf.function构建高效性能的Python代码。

一、tf.function是将函数编译为可调用的TensorFlow图。
1.tensorflow计算图

• 计算图是对有向图的表示，主要包含点和边；tensorflow使用计算图计算，计算图的点对应于ops（计算机点），variables（存储节点），constant，placeholder（1.0版本，数据节点）等；
• 边对应于Tensors，计算图中的边是有向边，定义了操作之间的关系，分为两类：一类用来传输数据，称为数据边；另一类用来定义依赖关系，称为控制边。所有的节点都通过数据边或者控制边连接，其中入度为0的节点没有前置依赖，可以立即执行；入度大于0的节点，要等待其依赖的所有节点执行结束之后，才可以执行。因此tensorflow主要包含两个部分：构建计算图和runtime运行计算图。

2.为什么要用计算图？

• 并行化，因为计算图是对计算的一种抽象，点之间的关系取决其依赖关系。因此，互相不依赖的计算可以并行计算，在多集群环境下可以进行分布式计算。
• 可移植性性，因为图是一种语言无关的表示方式，tensorflow中使用protobuf来存储图，可以使用C++，python，jave等语言来解析图。
• 将python的函数编译成tensorflow的图结构，所以可以方便的直接使用python的语法去写深度学习的一些函数代码
• 易于将模型导出为GraphDef+checkpoint和saveModel
• 使得eager exection默认打开。如果不使用@tf.function，虽然可以使用eager exection去写代码，但是模型的中间结果无法保存，所以无法进行预测
• tf1.0的代码可以经过@tf.function修饰之后在tf2里面继续使用。 此时其功能就是tf1里面的session

tf.function中的一个参数是autograph，默认是True，意思是在构建Graph时将自动使用AutoGraph，这样你可以在函数内部使用Python原生的条件判断以及循环语句，因为它们会被tf.cond和tf.while_loop转化为Graph代码。注意的一点是判断分支和循环必须依赖于Tensors才会被转化，当autograph为False时，如果存在判断分支和循环必须依赖于Tensors的情况将会出错。很容易理解，应用tf.function之后是Graph模式，Tensors是不能被遍历的，但是采用AutoGraph可以将其转换为Graph代码，所以可以成功。

tf.function(
func=None, input_signature=None,autograph=True,experimental_implements=None,
experimental_autograph_options=None, experimental_relax_shapes=False,
experimental_compile=None
)

二、Autograph将普通Python转换为TensorFlow图形代码。

• AutoGraph在TensorFlow1.x已经推出，主要是可以将一些常用的Python代码转化为TensorFlow支持的Graph代码。一个典型的例子是在TensorFlow中我们必须使用tf.while和tf.cond等复杂的算子来实现动态流程控制，但是现在我们可以使用Python原生的for和if等语法写代码，然后采用AutoGraph转化为TensorFlow所支持的代码。
• 需要注意的是，不是所有的函数都可以通过tf.function进行加速的。有的任务并不值得将函数转化为计算图形式，比如简单的矩阵乘法。然而，对于大量的计算，如对深度神经网络的优化，这一图转换能给性能带来巨大的提升。我们也把这样的图转化叫作tf.AutoGraph．在Tensorflow2.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）对于任何一个对该函数的调用，都会利用缓存中的计算图进行计算。

这里有一个注意点，在用@tf.function 进行修饰函数时，需要将申明的变量放在函数体的外面否则会报错。
原因：
这是因为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装饰。

代码示例：

import tensorflow as tf 
import numpy as  np 

#定义一个函数，将其转换为图结构并转换为tensorflow图代码
def scaled_elu(z, scale=1.0, alpha=1.0):
    # z >= 0 ? scale * z : scale * alpha * tf.nn.elu(z)
    #tf.nn.elu ：exp(z) - 1 if < 0, z otherwise.
    is_positive = tf.greater_equal(z, 0.0) # 返回元素的真值(x >= y)
    return scale * tf.where(is_positive, z, alpha * tf.nn.elu(z))

print(scaled_elu(tf.constant(-3.)))
print(scaled_elu(tf.constant([-3., 2.0])))
print("\n") 

#将函数转换为图结构 : tf.function()
scaled_elu_tf = tf.function(scaled_elu)
print(scaled_elu_tf(tf.constant(-3.)))
print(scaled_elu_tf(tf.constant([-3., 2.0])))

tf.Tensor(-0.95021296, shape=(), dtype=float32)
tf.Tensor([-0.95021296 2. ], shape=(2,), dtype=float32)

tf.Tensor(-0.95021296, shape=(), dtype=float32)
tf.Tensor([-0.950212962. ], shape=(2,), dtype=float32)

#使用python_function 可以查看图结构原生态的python函数
print(scaled_elu_tf.python_function is scaled_elu)

True

#查看转为TensorFlow图结构时，运行速度的变化
%timeit scaled_elu(tf.random.normal((1000, 1000)))
%timeit scaled_elu_tf(tf.random.normal((1000, 1000)))

2.82 ms ± 181 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
2.46 ms ± 50.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

#添加修饰符，转换为TensorFlow图结构

# 1 + 1/2 + 1/2^2 + ... + 1/2^n

@tf.function 
def converge_to_2(n_iters):
    total = tf.constant(0.)
    increment = tf.constant(1.0)
    for _ in  range(n_iters):
        total += increment
        increment /= 2.0
    return total
#这样我们就将上述函数转化为图结构了

print(converge_to_2(20))

tf.Tensor(1.9999981, shape=(), dtype=float32)

#使用autograph将普通Python转换为TensorFlow图形代码

def display_tf_code(func):
    code = tf.autograph.to_code(func)
    #用Markdown展示代码
    from IPython.display import display,Markdown
    display(Markdown('```python\n{}\n```'.format(code)))

#转化函数为图结构
display_tf_code(scaled_elu)

结果：

def tf__scaled_elu(z, scale=None, alpha=None):
  do_return = False
  retval_ = ag__.UndefinedReturnValue()
  with ag__.FunctionScope('scaled_elu', 'scaled_elu_scope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as scaled_elu_scope:
    is_positive = ag__.converted_call(tf.greater_equal, scaled_elu_scope.callopts, (z, 0.0), None, scaled_elu_scope)
    do_return = True
    retval_ = scaled_elu_scope.mark_return_value(scale * ag__.converted_call(tf.where, scaled_elu_scope.callopts, (is_positive, z, alpha * ag__.converted_call(tf.nn.elu, scaled_elu_scope.callopts, (z,), None, scaled_elu_scope)), None, scaled_elu_scope))
  do_return,
  return ag__.retval(retval_)

display_tf_code(converge_to_2)
#由于在使用@tf.function时默认autograph=True,这时coverage_to_2已经被转化为TensorFlow代码在进行转化是会报错

ConversionError: converting : AttributeError: ‘Function’ object has no attribute ‘code’

# 添加function修饰符时，申明变量要在函数体外，否则会报错

@tf.function
def add_21():
    var = tf.Variable(0.)
    return var.assign(21)

print(add_21())

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

@tf.function
def add_21():
    return var.assign(21)

var = tf.Variable(0.)
print(add_21())

tf.Tensor(21.0, shape=(), dtype=float32)