使用 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 确实能给我们性能带来很大的提升。