TF2默认的即时执行模式(Eager Execution)带来了灵活及易调试的特性,但在特定的场合,例如追求高性能或部署模型时,我们依然希望使用 TensorFlow 1.X 中默认的图执行模式(Graph Execution),将模型转换为高效的 TensorFlow 图模型。此时,TensorFlow 2 为我们提供了 tf.function
模块,结合 AutoGraph 机制,使得我们仅需加入一个简单的 @tf.function
修饰符,就能轻松将模型以图执行模式运行。
tf.function
基础使用方法
在 TensorFlow 2 中,推荐使用 tf.function
(而非 1.X 中的 tf.Session
)实现图执行模式,从而将模型转换为易于部署且高性能的 TensorFlow 图模型。只需要将我们希望以图执行模式运行的代码封装在一个函数内,并在函数前加上 @tf.function
即可,如下例所示。关于图执行模式的深入探讨可参考 附录 。
警告!
并不是任何函数都可以被 @tf.function
修饰!@tf.function
使用静态编译将函数内的代码转换成计算图,因此对函数内可使用的语句有一定限制(仅支持 Python 语言的一个子集),且需要函数内的操作本身能够被构建为计算图。建议在函数内只使用 TensorFlow 的原生操作,不要使用过于复杂的 Python 语句,函数参数只包括 TensorFlow 张量或 NumPy 数组,并最好是能够按照计算图的思想去构建函数(换言之,@tf.function
只是给了你一种更方便的写计算图的方法,而不是一颗能给任何函数加速的 银子弹 )。详细内容可参考 AutoGraph Capabilities and Limitations 。建议配合 附录 一同阅读本节以获得较深入的理解。
1 import tensorflow as tf 2 import time 3 from zh.model.mnist.cnn import CNN 4 from zh.model.utils import MNISTLoader 5 6 num_batches = 1000 7 batch_size = 50 8 learning_rate = 0.001 9 data_loader = MNISTLoader() 10 model = CNN() 11 optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate) 12 13 @tf.function 14 def train_one_step(X, y): 15 with tf.GradientTape() as tape: 16 y_pred = model(X) 17 loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred) 18 loss = tf.reduce_mean(loss) 19 # 注意这里使用了TensorFlow内置的tf.print()。@tf.function不支持Python内置的print方法 20 tf.print("loss", loss) 21 grads = tape.gradient(loss, model.variables) 22 optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables)) 23 24 start_time = time.time() 25 for batch_index in range(num_batches): 26 X, y = data_loader.get_batch(batch_size) 27 train_one_step(X, y) 28 end_time = time.time() 29 print(end_time - start_time)
运行 400 个 Batch 进行测试,加入 @tf.function
的程序耗时 35.5 秒,未加入 @tf.function
的纯即时执行模式程序耗时 43.8 秒。可见 @tf.function
带来了一定的性能提升。一般而言,当模型由较多小的操作组成的时候, @tf.function
带来的提升效果较大。而当模型的操作数量较少,但单一操作均很耗时的时候,则 @tf.function
带来的性能提升不会太大。
tf.function
内在机制
当被 @tf.function
修饰的函数第一次被调用的时候,进行以下操作:
-
在即时执行模式关闭的环境下,函数内的代码依次运行。也就是说,每个
tf.
方法都只是定义了计算节点,而并没有进行任何实质的计算。这与 TensorFlow 1.X 的图执行模式是一致的; -
使用 AutoGraph 将函数中的 Python 控制流语句转换成 TensorFlow 计算图中的对应节点(比如说
while
和for
语句转换为tf.while
,if
语句转换为tf.cond
等等; -
基于上面的两步,建立函数内代码的计算图表示(为了保证图的计算顺序,图中还会自动加入一些
tf.control_dependencies
节点); -
运行一次这个计算图;
-
基于函数的名字和输入的函数参数的类型生成一个哈希值,并将建立的计算图缓存到一个哈希表中。
在被 @tf.function
修饰的函数之后再次被调用的时候,根据函数名和输入的函数参数的类型计算哈希值,检查哈希表中是否已经有了对应计算图的缓存。如果是,则直接使用已缓存的计算图,否则重新按上述步骤建立计算图。
提示!
对于熟悉 TensorFlow 1.X 的开发者,如果想要直接获得 tf.function
所生成的计算图以进行进一步处理和调试,可以使用被修饰函数的 get_concrete_function
方法。该方法接受的参数与被修饰函数相同。例如,为了获取前节被 @tf.function
修饰的函数 train_one_step
所生成的计算图,可以使用以下代码:
graph = train_one_step.get_concrete_function(X, y)
其中 graph
即为一个 tf.Graph
对象。
以下是一个测试题:
import tensorflow as tf import numpy as np @tf.function def f(x): print("The function is running in Python") tf.print(x) a = tf.constant(1, dtype=tf.int32) f(a) b = tf.constant(2, dtype=tf.int32) f(b) b_ = np.array(2, dtype=np.int32) f(b_) c = tf.constant(0.1, dtype=tf.float32) f(c) d = tf.constant(0.2, dtype=tf.float32) f(d)
思考一下,上面这段程序的结果是什么?答案是:
The function is running in Python 1 2 2 The function is running in Python 0.1 0.2
当计算 f(a)
时,由于是第一次调用该函数,TensorFlow 进行了以下操作:
-
将函数内的代码依次运行了一遍(因此输出了文本);
-
构建了计算图,然后运行了一次该计算图(因此输出了 1)。这里
tf.print(x)
可以作为计算图的节点,但 Python 内置的print
则不能被转换成计算图的节点。因此,计算图中只包含了tf.print(x)
这一操作; -
将该计算图缓存到了一个哈希表中(如果之后再有类型为
tf.int32
,shape 为空的张量输入,则重复使用已构建的计算图)。
计算 f(b)
时,由于 b 的类型与 a 相同,所以 TensorFlow 重复使用了之前已构建的计算图并运行(因此输出了 2)。这里由于并没有真正地逐行运行函数中的代码,所以函数第一行的文本输出代码没有运行。计算 f(b_)
时,TensorFlow 自动将 numpy 的数据结构转换成了 TensorFlow 中的张量,因此依然能够复用之前已构建的计算图。
计算 f(c)
时,虽然张量 c
的 shape 和 a
、 b
均相同,但类型为 tf.float32
,因此 TensorFlow 重新运行了函数内代码(从而再次输出了文本)并建立了一个输入为 tf.float32
类型的计算图。
计算 f(d)
时,由于 d
和 c
的类型相同,所以 TensorFlow 复用了计算图,同理没有输出文本。
而对于 @tf.function
对 Python 内置的整数和浮点数类型的处理方式,我们通过以下示例展现:
f(d) f(1) f(2) f(1) f(0.1) f(0.2) f(0.1)
结果为:
The function is running in Python 1 The function is running in Python 2 1 The function is running in Python 0.1 The function is running in Python 0.2 0.1
简而言之,对于 Python 内置的整数和浮点数类型,只有当值完全一致的时候, @tf.function
才会复用之前建立的计算图,而并不会自动将 Python 内置的整数或浮点数等转换成张量。因此,当函数参数包含 Python 内置整数或浮点数时,需要格外小心。一般而言,应当只在指定超参数等少数场合使用 Python 内置类型作为被 @tf.function
修饰的函数的参数。
下一个思考题:
import tensorflow as tf a = tf.Variable(0.0) @tf.function def g(): a.assign(a + 1.0) return a print(g()) print(g()) print(g())
这段代码的输出是:
tf.Tensor(1.0, shape=(), dtype=float32)
tf.Tensor(2.0, shape=(), dtype=float32)
tf.Tensor(3.0, shape=(), dtype=float32)
正如同正文里的例子一样,你可以在被 @tf.function
修饰的函数里调用 tf.Variable
、 tf.keras.optimizers
、 tf.keras.Model
等包含有变量的数据结构。一旦被调用,这些结构将作为隐含的参数提供给函数。当这些结构内的值在函数内被修改时,在函数外也同样生效。
AutoGraph:将 Python 控制流转换为 TensorFlow 计算图
前面提到,@tf.function
使用名为 AutoGraph 的机制将函数中的 Python 控制流语句转换成 TensorFlow 计算图中的对应节点。以下是一个示例,使用 tf.autograph
模块的低层 API tf.autograph.to_code
将函数 square_if_positive
转换成 TensorFlow 计算图:
import tensorflow as tf @tf.function def square_if_positive(x): if x > 0: x = x * x else: x = 0 return x a = tf.constant(1) b = tf.constant(-1) print(square_if_positive(a), square_if_positive(b)) print(tf.autograph.to_code(square_if_positive.python_function))
输出:
tf.Tensor(1, shape=(), dtype=int32) tf.Tensor(0, shape=(), dtype=int32) def tf__square_if_positive(x): do_return = False retval_ = ag__.UndefinedReturnValue() cond = x > 0 def get_state(): return () def set_state(_): pass def if_true(): x_1, = x, x_1 = x_1 * x_1 return x_1 def if_false(): x = 0 return x x = ag__.if_stmt(cond, if_true, if_false, get_state, set_state) do_return = True retval_ = x cond_1 = ag__.is_undefined_return(retval_) def get_state_1(): return () def set_state_1(_): pass def if_true_1(): retval_ = None return retval_ def if_false_1(): return retval_ retval_ = ag__.if_stmt(cond_1, if_true_1, if_false_1, get_state_1, set_state_1) return retval_
我们注意到,原函数中的 Python 控制流 if...else...
被转换为了 x = ag__.if_stmt(cond, if_true, if_false, get_state, set_state)
这种计算图式的写法。AutoGraph 起到了类似编译器的作用,能够帮助我们通过更加自然的 Python 控制流轻松地构建带有条件 / 循环的计算图,而无需手动使用 TensorFlow 的 API 进行构建。
使用传统的 tf.Session
不过,如果你依然钟情于 TensorFlow 传统的图执行模式也没有问题。TensorFlow 2 提供了 tf.compat.v1
模块以支持 TensorFlow 1.X 版本的 API。同时,只要在编写模型的时候稍加注意,Keras 的模型是可以同时兼容即时执行模式和图执行模式的。注意,在图执行模式下, model(input_tensor)
只需运行一次以完成图的建立操作。
例如,通过以下代码,同样可以在 MNIST 数据集上训练前面所建立的 MLP 或 CNN 模型:
1 optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=learning_rate) 2 num_batches = int(data_loader.num_train_data // batch_size * num_epochs) 3 # 建立计算图 4 X_placeholder = tf.compat.v1.placeholder(name='X', shape=[None, 28, 28, 1], dtype=tf.float32) 5 y_placeholder = tf.compat.v1.placeholder(name='y', shape=[None], dtype=tf.int32) 6 y_pred = model(X_placeholder) 7 loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y_placeholder, y_pred=y_pred) 8 loss = tf.reduce_mean(loss) 9 train_op = optimizer.minimize(loss) 10 sparse_categorical_accuracy = tf.keras.metrics.SparseCategoricalAccuracy() 11 # 建立Session 12 with tf.compat.v1.Session() as sess: 13 sess.run(tf.compat.v1.global_variables_initializer()) 14 for batch_index in range(num_batches): 15 X, y = data_loader.get_batch(batch_size) 16 # 使用Session.run()将数据送入计算图节点,进行训练以及计算损失函数 17 _, loss_value = sess.run([train_op, loss], feed_dict={X_placeholder: X, y_placeholder: y}) 18 print("batch %d: loss %f" % (batch_index, loss_value)) 19 20 num_batches = int(data_loader.num_test_data // batch_size) 21 for batch_index in range(num_batches): 22 start_index, end_index = batch_index * batch_size, (batch_index + 1) * batch_size 23 y_pred = model.predict(data_loader.test_data[start_index: end_index]) 24 sess.run(sparse_categorical_accuracy.update(y_true=data_loader.test_label[start_index: end_index], y_pred=y_pred)) 25 print("test accuracy: %f" % sess.run(sparse_categorical_accuracy.result()))