自定义
使用tf.function时的性能
Tensorflow 2.0中一个主要的改变就是移除tf.Session这一概念.这样可以帮助用户更好的组织代码,不用将tf.Session作为一个变量在Python函数中传来传去,我们可以用一个Python装饰符来进行加速,那就是@tf.function。
需要注意的是,不是所有的函数都可以通过tf.function进行加速的。有的任务并不值得将函数转化为计算图形式,比如简单的矩阵乘法。然而,对于大量的计算,如对深度神经网络的优化,这一图转换能给性能带来巨大的提升。我们也把这样的图转化叫作tf.AutoGraph。在Tensorflow 2.0中,我们会自动的对被@tf.function装饰的函数进行AutoGraph优化。
我们来粗浅的看一下被tf.function装饰的函数第一次执行时都做了什么:
- 1、函数被执行并且被跟踪(tracing)。Eager execution处于关闭状态,所有的Tensorflow函数被当做tf.Operation进行图的创建。
- 2、AutoGraph被唤醒,去检测Python代码可以转为Tensorflow的逻辑,比如while > tf.while, for > tf.while, if > tf.cond, assert > tf.assert。
- 3、通过以上两步,我们对函数进行建图,为了保证Python代码中每一行的执行顺序,tf.control_dependencies被自动加入到代码中。保证第i行执行完后我们会执行第i+1行。
- 4、返回tf.Graph,根据函数名和输入参数,我们将这个graph存到一个cache中。
- 5、对于任何一个该函数的调用,我们会重复利用cache中的计算图进行计算。