使用 tf.function 提升效率

在之前的入门介绍之中,我们曾经介绍过 TensorFlow1.x 采用的并不是 Eager execution 执行模型;而 TensorFlow2.x 默认采用的是 Eager execution 模式。

这种改变使得我们可以更加容易地学习,但是也会造成性能的损失,因此, TensorFlow 在 2.0 版本之后引入了 tf.function 。

1. 什么是 tf.funtion

在 TensorFlow1.x 之中,如果我们想要运行一个学习任务,那么我们需要首先创建一个 tf.Sesstion (),然后再调用 Session.run () 进行运行。

其实在 TensorFlow1.x 内部,当我们在 TensorFlow 之中进行工作的时候, TensorFlow 会帮助我们创建一个计算图 tf.graph ,然后通过 tf.Session 对计算图进行计算

而在 TensorFlow2.x 之中,其默认采用的是 Eager execution 执行方式,在该执行方式之中,我们不再需要定义一个计算图来进行

这样就产生了一些问题:

  • 使用 tf.Sesstion () 的运行效率非常高,但是代码很难懂
  • 使用 Eager execution 方式的代码很简单,但是执行效率比较低

有什么方法能够兼顾两者吗?

那就是 tf.function 。

2. tf.funtion 的用法

tf.function 是一个函数标注修饰,也就是如下的形式:

@tf.function
def my_function():
    ...

其实如你所见,这就是 tf.function 的全部用法。

我们只需要在我们要修饰的函数之前加上 tf.function 标注既可

采用 tf.function,TensorFlow 会将该函数转变为计算图 tf.graph 的形式来进行运算,这会使得该函数在进行大量运算的时候会加速非常多。

是不是所有的函数都适合 tf.function 进行修饰呢?

答案是否定的,以下两种情况不适合使用 tf.function 进行修饰:

  • 函数本身的计算非常简单,那么构建计算图本身的时间就会相对非常浪费;
  • 当我们需要在函数之中定义 tf.Variable 的时候,因为 tf.function 可能会被调用多次,因此定义 tf.Variable 会产生重复定义的情况。

3. tf.function 的性能

既然了解了 tf.function 的用法,那么我们便来测试一下 tf.function 的性能,我们采用一个简单的卷积神经网络来进行测试:

import tensorflow as tf
import timeit


def f1(layer, image):
    y = layer(image)
    return y

@tf.function
def f2(layer, image):
    y = layer(image)
    return y

layer = tf.keras.layers.Conv2D(300, 3)
image = tf.zeros([64, 32, 32, 3])

model = tf.keras.models.Sequential([
        tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(128, (3, 3), activation='relu'),  
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')])

print(timeit.timeit(lambda: f1(model, image), number=500))
print(timeit.timeit(lambda: f2(model, image), number=500))

在这里,我们定义了两个相同的函数,其中一个使用了 tf.function 进行修饰,而另外一个没有。

在这里我们使用 lambda 函数来让函数重复执行 500 次,并且使用 timeit 来进行时间的统计,得到两个函数的执行时间,从而进行比较。

最终,我们可以得到结果:

17.20403664399987
12.07886406200032

由此可以看出,我们的 tf.function 已经提升了一定的速度,但是提升的速度有限,目前大概提升了 25 % 的速度。这是因为我们的计算仍然还是太简单了,当我们计算非常大的时候,性能会有很大的提升。

4. 小结

在这节课之中,我们学习到了什么是 tf.function ,以及 tf.function 的基本原理,然后我们了解了 tf.function 的使用方法;最后我们通过一个简单的神经网络来进行了性能的测试,最终我们发现我们的 tf.function 确实能给我们性能带来很大的提升。