Press "Enter" to skip to content

Tensorflow 2.0上手6: 解剖tf.function的使用

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

Session execution

 

了解Tensorflow 1.x的读者都知道Tensorflow一般的工作流程.创建一个计算图 tf.graph
,然后通过 tf.Session
对计算图进行计算.以下是一段简单的代码:

 

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

 

用Tensorflow 2.0默认的Eager execution实现则很不一样,用户不再需要直接定义计算图或者通过 tf.Session
来执行代码,也不需要调用 tf.global_variables_initializer
去初始化变量或者通过 tf.control_dependencies
去执行计算图中没有包含的节点.

 

Tensorflow变得像普通的Python代码一样简单.

 

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的解释器来进行计算,无法对数据流以及计算图进行优化.

 

为了在代码可读性和代码速度之间保持平衡,Tensorflow 2.0引入了 tf.function
这一概念.

 

tf.function, not tf.Session

 

Tensorflow 2.0中一个主要的改变就是移除 tf.Session
这一概念.这样可以帮助用户更好的组织代码,不用将 tf.Session
作为一个变量在Python函数中传来传去,我们可以用一个Python装饰符来进行加速,那就是 @tf.function

 

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

 

懒人介绍tf.function

 

我们来粗浅的看一下被 tf.function
装饰的函数第一次执行时都做了什幺:

 

tf.Operation
tf.control_dependencies

 

我们来看一下Tensorflow 2.0中Eager Execution的代码如何转为tf.function的代码.首先来看一段简单的Tensorflow 2.0代码:

 

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 yprint(f().numpy())#执行结果
[[22. 22.]
 [23. 13.]]

 

因为Tensorflow 2.0默认是Eager execution,代码的阅读和执行就和普通的Python代码一样,简单易读.

 

From eager to tf.function

 

首先我们简单的加上 @tf.function
装饰一下,为了方便调试,我们加入一个 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 yf()#执行结果
PRINT: Tensor("add:0", shape=(2, 2), dtype=float32)
ValueError: tf.function-decorated function tried to create variables on non-first call.

 

咦,为什幺代码报错了?

 

我以为我发现了一个Bug,因为我理想的执行方式是第一次执行后进行构图并且调用这个图进行计算,之后每次都重复调用这个计算图,然而我们却收到了异常.

 

但实际上程序的执行过程和我们期待的是一样的,因为 tf.function
可能会对一段Python函数进行多次执行来构图,在多次执行的过程中,同样的Variable被创建了多次,产生错误.

 

这其实也是一个很容易混乱的概念,在eager mode下一个Variable是一个Python object,所以会在执行范围外被销毁.但是在 tf.function
的装饰下,Variable变成了 tf.Variable
,是在Graph中持续存在的.

 

把一个在eager mode下正常执行的函数转换到Tensorflow图形式,需要一边思考着计算图一边构建程序.

 

所以我们在使用 tf.function
有几种操作需要做:

 

 

    1. 设计函数f时需要一些输入参数,这个输入参数可以是

tf.Variable

    1. 或者其他任何类型.

 

    1. 设计一个函数从parent scope继承Python variable,在函数中检查Variable是否已经定义过

(if b != None)

 

    1. 将所有的内容写到一个class里,就好像Keras layer一样,所有的Variable都是class的内部参数

(self._b)

    1. ,将class的

__call__()

    1. 通过

tf.function

    1. 装饰.

 

 

我们来分析一下上面这些写法.

 

方案2 vs 方案3

 

上面的第2种方案和第3种方案类似,但是从面向对象的角度来看,方案3会更美一些.

 

我们先来看一下方案2的丑陋写法(很不推荐):

 

b = [email protected]
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 yf()

 

方案3的写法:

 

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

 

函数逻辑与主过程的分割相对清晰,没有全局变量,class F可以被实例化然后调用,而不用担心我们生成一个全局变量b,让其他函数也能看到.

 

运用上述两种代码,我们已经解决了对于函数创造自己的Variable的问题.在实际运行过程中, tf.function
生成的函数将和eager execution返回相同的结果.

 

方案1举例

 

现在来看一下上面讲到的方案1,将变量传递到函数中去.

 

@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 yb = tf.Variable(12.)
f(b)

 

如上一节所述,该函数也可以产生预期的行为。此外,通过传递Variable,我们也可以在函数内部更新Variable的值.比方下面的代码可以更新参数x,产生1,2,3.

 

@tf.function
def g(x):
 x.assign_add(1)
 return xa = tf.Variable(0)
print(g(a))
print(g(a))
print(g(a))#执行结果
tf.Tensor(1, shape=(), dtype=int32)
tf.Tensor(2, shape=(), dtype=int32)
tf.Tensor(3, shape=(), dtype=int32)

 

Conclusions

 

关于 tf.function
今天就先讲到这里,我们大致了解了怎幺把一个Python eager mode下执行的函数通过 tf.function
转化为更快的计算图.

 

等下一章我们来仔细研究一下当我们把一个 tf.Tensor
或者Python value传入被 tf.function
装饰的代码里会发生什幺.以及 tf.function
第一次被执行时,里面的Python代码究竟是如何运行的,每一步都会被转为我们期待的计算图形式吗?

 

原文:
Dissecting tf function part 1 https://pgaleone.eu/tensorflow/tf.function/2019/03/21/dissecting-tf-function-part-1/

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注