一、什么是Autograph
在前一篇文章TensorFlow核心概念之计算图中我们提到过,TensorFlow中的构建方式主要有三种,分别是:静态计算图构建、动态计算图构建和Autograph。其中静态计算图主要是在TensorFlow1.0中支持的计算图构建方式,这种方式构建的计算图虽然执行效率高,但不便于编码过程中的调试,交互体验差。因此2.0之后TensorFlow开始支持动态计算图,虽然便于了编码过程中调试和交互体验,但是执行效率问题随之而来。于是就有了Autograph,Autograph是一种将动态图转换成静态图的实现机制,通过在普通python方法上使用@tf.function
进行装饰,从而将动态图转换成静态图。
三、Autograph实现原理
为了搞清楚Autograph的机制原理,我们需要知道,当我们使用@tf.function
装饰一个函数后,在调用这些函数时,TensorFlow到底做了什么?下面我们详细介绍Autograph的实现原理。当调用被@tf.function
时,TensorFlow一共做了两件事:第一件事是创建静态计算图,第二件事是执行静态计算图。执行计算图没什么好讲的,就是针对创建好的计算图,根据输入的参数进行执行,关键的问题是TensorFlow是如何创建计算这个静态计算图的。
当执行被@tf.function
装饰的函数时,TensorFlow会在后端隐式的创建一个静态计算图,静态计算图的创建过程大体时这样的:跟踪执行一遍函数体中的Python代码,确定各个变量的Tensor类型,并根据执行顺序将各TensorFlow的算子添加到计算图中。 在该过程中,如果@tf.function(autograph=True)
(默认开启autograph),TensorFlow会将Python控制流转换成TensorFlow的静态图控制流。 主要是将if语句转换成 tf.cond算子表达,将while和for循环语句转换成tf.while_loop算子表达,并在必要的时候添加 tf.control_dependencies指定执行顺序依赖关系。这里需要注意的是,非TensorFlow的函数不会被添加到计算图中,也就是说,像Python原生支持的一些函数在构建静态计算图的过程中,只会被跟踪执行,不会将该函数作为算子嵌入到TensorFlow的静态计算图中。
另外还需要注意的一点是,当在调用@tf.function装饰的函数时,如果输入的参数是Tensor类型,此时TensorFlow会从性能的角度出发,去判断当前入参类型下的静态计算图是否已经存在,如果已经存在,则直接执行计算图,从而省去构建静态计算图的过程,进而提升效率。但是如果发现当前入参的静态计算图不存在,则需要重新创建新的计算图。另外需要注意的是,如果调用被@tf.function装饰的函数时,入参不是Tensor类型,则每次调用的时候都需要先创建静态计算图,然后执行计算图。
三、Autograph的编码规范
介绍完TensorFlow的实现原理,下面我们简单介绍一下Autograph的编码规范和使用建议。并通过简单的示例来演示为什么要有这些规范和建议。
1. 被@tf.function
修饰的函数应尽量使用TensorFlow中的函数,而非外部函数。
2. 不能在@tf.function
修饰的函数内部定义tf.Variable变量。
3. 被@tf.function
修饰的函数不可修改该函数外部的Python列表或字典等数据结构变量。
4. 调用被@tf.function
修饰的函数,入参尽量使用Tensor类型。
四、Autograph的编码规范解析
1. 被@tf.function
修饰的函数应尽量使用TensorFlow中的函数,而非外部函数。
我们可以看下面一段代码,我们定义了两个@tf.function
修饰的函数,其中第一个函数体内使用了两个外部函数,分别是np.random.randn(3,3)
和print('---------')
,第二个函数体内全部使用TensorFlow中的函数。
import numpy as np
import tensorflow as tf
@tf.function
def np_random():
a = np.random.randn(3,3)
tf.print(a)
print('---------')
@tf.function
def tf_random():
a = tf.random.normal((3,3))
tf.print(a)
tf.print('---------')
下面我们调用两次第一个被@tf.function
修饰的函数:
print('第1次调用:')
np_random()
print('第2次调用:')
np_random()
结果如下:
第1次调用:
---------
array([[ 0.78826988, -0.05816027, 0.88905733],
[-1.98118034, -0.10032147, -0.51427141],
[ 0.50533615, -1.11163988, -0.87748809]])
第2次调用:
array([[ 0.78826988, -0.05816027, 0.88905733],
[-1.98118034, -0.10032147, -0.51427141],
[ 0.50533615, -1.11163988, -0.87748809]])
这个时候我们会发现三个问题:
- 第一次调用的时候,
print('---------')
方法执行了,最起码看起是执行了,也确实是执行了,而第二次调用的时候,print('---------')
方法并没有执行; - 第一次调用的时候,
print('---------')
方法在tf.print(a)
之前调用了; - 两次调用之后,变量
a
的结果是一样的。
下面针对以上问题,我们来详细解释一下:首先在第一次调用的是,会进行静态计算图的创建,这个时候Python后端会跟踪执行一遍函数体Python的代码,,并将方法体中的变量和算子进行映射和加入计算图中,这里需要注意的是,由于np.random.randn(3,3)
和print('---------')
方法并不是TensorFlow中的方法,因此无法加入到计算图中,因此只有tf.print(a)
方法加入到了静态计算图中,因此只有在第一次创建计算图的时候进行跟踪执行,而第二次执行时,如果计算图已经存在,这个时候时不需要再执行的,这也就是为什么print('---------')
会先在tf.print(a)
前面执行,且执行一次。因为在实际执行计算图的过程中,都只会执行tf.print(a)
这一个方法,这也导致了为什么多次调用之后,打印出来的a
的结果是一样的。基于以上原因,我们再两次调用一下第二个方法tf_random()
,示例代码和结果如下:
print('第1次调用:')
tf_random()
print('第2次调用:')
tf_random()
结果如下:
第1次调用:
[[1.47568643 -0.204902112 0.694708228]
[-0.868299544 1.65556359 0.520012081]
[-0.215179399 -0.400003046 -0.393970907]]
---------
第2次调用:
[[0.0756372586 1.06571424 -0.579676867]
[-0.937381923 -2.79628611 -1.38038337]
[-0.762175 -1.79867613 0.329570293]]
---------
这个时候我们可以看出,全部使用TensorFlow函数的方法调用的结果是符合我们的预期的。
2. 不能在@tf.function
修饰的函数内部定义tf.Variable变量。
这个我们就直接示例,代码如下:
@tf.function
def inner_var():
x = tf.Variable(1.0,dtype = tf.float32)
x.assign_add(1.0)
tf.print(x)
return(x)
这个时候执行的时候,代码会直接报错,报错信息如下:
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-12-c95a7c3c1ddd> in <module>
7
8 #执行将报错
----> 9 inner_var()
10 inner_var()
~/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py in __call__(self, *args, **kwds)
566 xla_context.Exit()
567 else:
--> 568 result = self._call(*args, **kwds)
569
570 if tracing_count == self._get_tracing_count():
......
ValueError: tf.function-decorated function tried to create variables on non-first call.
如果我们将这个变量拿到@tf.function
修饰的函数外,则可以直接执行,代码如下:
x = tf.Variable(1.0,dtype=tf.float32)
@tf.function
def outer_var():
x.assign_add(1.0)
tf.print(x)
return(x)
outer_var()
outer_var()
结果如下:
2
3
3. 被@tf.function
修饰的函数不可修改该函数外部的Python列表或字典等数据结构变量。
正对这个我们直接看代码示例,首先我们在不用@tf.function
修饰的函数来演示一下执行结果,代码如下:
tensor_list = []
def append_tensor(x):
tensor_list.append(x)
return tensor_list
append_tensor(tf.constant(1.0))
append_tensor(tf.constant(2.0))
print(tensor_list)
结果如下:
[<tf.Tensor: shape=(), dtype=float32, numpy=1.0>, <tf.Tensor: shape=(), dtype=float32, numpy=2.0>]
这个时候我们发现一切如我们的预期,没有任何问题,接下来我们对这个append_tensor(x)
函数加上@tf.function
修饰,代码如下:
tensor_list = []
@tf.function
def append_tensor(x):
tensor_list.append(x)
return tensor_list
append_tensor(tf.constant(1.0))
append_tensor(tf.constant(2.0))
print(tensor_list)
结果如下:
[<tf.Tensor 'x:0' shape=() dtype=float32>]
其实出现这个问题的原因呢也很好解释,那就是tensor_list.append(x)
不是一个TensorFlow的方法,在构建计算图的时候呢,这个方法并不会作为算子加入到静态计算图中,那么在最后执行计算图的时候,其实也就不会去执行这个方法了,这就是为啥最终这个列表内容为空的原因。
4. 调用被@tf.function
修饰的函数,入参尽量使用Tensor类型。
这一点是从性能的角度出发的,因为在调用被@tf.function
修饰的函数时,TensorFlow会根据入参类型来决定是否要重新创建静态计算图,这一点时从性能的角度出发的,对结果其实并没有实际的影响。示例代码如下:
import tensorflow as tf
import numpy as np
@tf.function(autograph=True)
def myadd(a,b):
c = a + b
print("tracing")#为了方便知道在创建计算图
tf.print(c)
return c
首先我们使用Tensor类型的入参多次调用该函数:
print("第1次调用:")
myadd(tf.constant("Hello"), tf.constant("World"))
print("第2次调用:")
myadd(tf.constant("Good"), tf.constant("Bye"))
结果如下:
第1次调用:
tracing
HelloWorld
第2次调用:
GoodBye
而当我们使用非Tensor类型的入参多次调用该函数:
print("第1次调用:")
myadd("Hello","World")
print("第2次调用:")
myadd("Good","Bye")
结果如下:
第1次调用:
tracing
HelloWorld
第2次调用:
tracing
GoodBye
这个时候我们发现,如果在调用@tf.function
修饰的函数时,如果入参的类型不是TensorFlow的类型,那么在多次调用该方法时,如果入参类型不变,内容变换的化,是需要多次创建静态计算图的,而如果使用Tensor类型的入参,则不会出现重复创建静态计算图的过程,除非入参类型改变,这样可以大大的提高调用性能。OK,关于TensorFlow中的Autograph就简单介绍这么多。