原 文:Effective TensorFlow 2
译 者:Xovee
许 可:文字:Creative Commons Attribution 4.0 License,代码:Apache 2.0 License
翻译时间:2019年12月20日
适用版本:TensorFlow 2.0
正文
在 TensorFlow 2.0 中有许多改动使得 TensorFlow 变得更加地富有生产力。TensorFlow 2.0 移除了许多冗余的 APIs,让 APIs 更一致(统一的 RNNs,同一的优化器),并且更好地与 Python 执行器以及与 Eager execution 结合。
许多 RFCs 已经解释了 TensorFlow 2.0 中的许多改动。本文给读者提供了一种 TensorFlow 2.0 的开发应该是怎样的视角。本文假设读者已经对 TensorFlow 1.x 有了一定的了解。
对 TensorFlow 2.0 改动的一个简要的总结
API 简化
许多 API 在 TF 2.0 中已经永久的消失了。其中一些主要的改动包括移除了tf.app
、tf.flags
,tf.logging
以支持新的开源软件 absl-py,对存在于tf.contrib
中的许多项目进行了搬家,并且简化了tf.*
中较少使用的函数(搬到次级包中例如tf.math
)。一些 APIs 被等同的模块替代了 - tf.summary
、tf.keras.metrics
、tf.keras.optimizers
。如果你想对这些重命名改动进行适配,最简单的方法是运行 v2 升级脚本。
Eager execution
TensorFlow 1.x 要求用户使用一种抽象语法树(图)来调用tf.*
APIs。它随后要求用户手动地编译抽象语法树,通过传递一系列输出 tensors 和输入 tensors 到session.run()
。在 TensorFlow 2.0 中,这些操作可以 eagerly 执行(就像普通的 Python 一样),并且在 2.0 中,图和 sessions 应当被看作是某种实现细节。
应用 eager execution 带来一个显著地副作用,即tf.control_dependecies()
不再需要了,因为所有的代码都会按顺序执行(与tf.function
一起)。
globals 也不再需要了
TensorFlow 1.x 极度依赖于隐形全局命名空间。当你调用tf.Variable()
,它会被放在默认图中,并且它会一直存在于那里,就算你丢失了指向它的 Python 变量。你可以随后恢复tf.Variable
,如果你知道它的名字并且它已经被创建了。但是如果你无法控制变量的创建的话,做到这一点会很困难。所以,所有的这些机制被用来帮助用户找到这些变量,以及帮助框架来找到用户创建的变量: Variable scopes,global collections,helper functions 例如tf.get_global_step()
、tf.global_variable_initializer()
,优化器隐式地计算可训练变量的梯度,等等等等。在 TensorFlow 2.0 中,所有的这些机制都被消除了(Variable 2.0 RFC),从而支持默认的机制:保持对变量的跟踪!如果你丢失了对tf.Variable
的跟踪,它会被垃圾回收(garbage collected)。
对变量追踪的需求会加大用户的负担,但是随着 Keras 对象的加入,这种负担被最小化了。
functions,而非 sessions
调用session.run()
就像调用一个普通的函数一样:你选取输入以及要调用的函数,然后你得到了一系列输出。在 TensorFlow 2.0 中,你可以对 Python 函数进行 tf.function
装饰,使得 TensorFlow 在运行它的时候把它看作是一个单独的图(Functions 2.0 RFC)。这个机制使得 TensorFlow 2.0 有了以下优点:
- 性能:函数得以被优化(node pruning,kernel fusion, 等等)
- 可移植性:函数可以被导出/重导入(SaveModel 2.0 RFC),允许用户去重用和分享 TensorFlow 函数。
# TensorFlow 1.X
outputs = session.run(f(placeholder), feed_dict={placeholder: input})
# TensorFlow 2.0
outputs = f(input)
用户可以随意地混合 Python 代码和 TensorFlow 代码来体现 Python 语言的表达性。但是可移植的 TensorFlow 可能在没有 Python 解释器的环境下执行,例如 mobile、C++,以及 JavaScript。为了帮助用户避免在使用@tf.function
的时候重构代码,AutoGraph自动地把 Python 结构转换为等价的 TensorFlow 语言:
for
/while
->tf.while_loop
(支持break
和continue
)if
->tf.cond
for _ in dataset
->dataset.reduce
AutoGraph 支持任意地控制流嵌套,从而可以精确地支持许多复杂的机器学习模型,例如 sequence model、强化学习、自定义训练循环,等等。
地道地 TensorFlow 2.0 代码
用轻量级的函数来重构你的代码
在 TensorFlow 1.x 中一个常用的做法是“洗碗槽”策略,那些可能会用到的计算,都会事先地定义好,然后使用session.run()
来执行它们。在 TensorFlow 2.0 中,用户应该使用轻量级的函数来重构他们的代码,只有在需要的时候才定义各种函数。一般来说,用户不需要对每一个轻量级的函数都用tf.function
来修饰;只有在定义高阶计算的时候(例如模型的每一步训练、正向传播等),才用tf.function
来修饰它。
使用 Keras layers 和 models 来管理变量
Keras 的 models 和 layers 提供了非常方便的variables
和trainable_variables
属性,它们可以递归地收集所有依赖的变量。这使得管理本地变量非常地方面。
对比:
def dense(x, W, b):
return tf.nn.sigmoid(tf.matmul(x, W) + b)
@tf.function
def multilayer_perceptron(x, w0, b0, w1, b1, w2, b2 ...):
x = dense(x, w0, b0)
x = dense(x, w1, b1)
x = dense(x, w2, b2)
...
# You still have to manage w_i and b_i, and their shapes are defined far away from the code.
使用 Keras:
# Each layer can be called, with a signature equivalent to linear(x)
layers = [tf.keras.layers.Dense(hidden_size, activation=tf.nn.sigmoid) for _ in range(n)]
perceptron = tf.keras.Sequential(layers)
# layers[3].trainable_variables => returns [w3, b3]
# perceptron.trainable_variables => returns [w0, b0, ...]
Keras 的 layers 和 models 继承自tf.train.Checkpointable
,并且与@tf.function
进行了整合,这使得这些 layers 和 models (它们是 Keras 对象) 可以直接利用检查点(checkpoint),或导出为 SavedModels。你并不总是需要利用 Keras 的.fit
API 来获取这些整合的优点。
下面是一个迁移学习的例子,来说明为什么 Keras 可以让收集相关变量变得非常地容易。假设我们现在需要训练一个 multi-headed model with a shared trunk:
trunk = tf.keras.Sequential([...])
head1 = tf.keras.Sequential([...])
head2 = tf.keras.Sequential([...])
path1 = tf.keras.Sequential([trunk, head1])
path2 = tf.keras.Sequential([trunk, head2])
# Train on primary dataset
for x, y in main_dataset:
with tf.GradientTape() as tape:
prediction = path1(x)
loss = loss_fn_head1(prediction, y)
# Simultaneously optimize trunk and head1 weights.
gradients = tape.gradient(loss, path1.trainable_variables)
optimizer.apply_gradients(zip(gradients, path1.trainable_variables))
# Fine-tune second head, reusing the trunk
for x, y in small_dataset:
with tf.GradientTape() as tape:
prediction = path2(x)
loss = loss_fn_head2(prediction, y)
# Only optimize head2 weights, not trunk weights
gradients = tape.gradient(loss, head2.trainable_variables)
optimizer.apply_gradients(zip(gradients, head2.trainable_variables))
# You can publish just the trunk computation for other people to reuse.
tf.saved_model.save(trunk, output_path)
利用 tf.data.Datasets
和 @tf.function
在对训练数据进行迭代的时候,你可以使用标准的 Python 循环。或者,使用tf.data.Dataset
,它是从磁盘中读取训练数据的最好的方法。数据集是 iterables(不是 iterators),它可以像其他的 Python iterables 一样在 Eager mode 环境下工作。你可以完全地利用数据集异步预读取功能,通过tf.function()
将 Python 迭代替代为等价的使用 AutoGraph 的操作。
@tf.function
def train(model, dataset, optimizer):
for x, y in dataset:
with tf.GradientTape() as tape:
prediction = model(x)
loss = loss_fn(prediction, y)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
如果你使用 Keras 的.fit()
API,你不需要担心数据集的迭代。
model.compile(optimizer=optimizer, loss=loss_fn)
model.fit(dataset)
在 Python 控制流中利用 AutoGraph 的优点
AutoGraph 提供了一种方法去转换数据依赖的控制流到基于图的控制流,例如tf.cnd
和tf.while_loop
。
一种经常用到数据依赖控制流的地方是在序列模型之中。tf.keras.layers.RNN
包装了一个 RNN cell,允许你去静态地或者动态地执行回归(unroll recurrence)。为了论证的方便,你可以重新实现一个动态的 unroll:
class DynamicRNN(tf.keras.Model):
def __init__(self, rnn_cell):
super(DynamicRNN, self).__init__(self)
self.cell = rnn_cell
def call(self, input_data):
# [batch, time, features] -> [time, batch, features]
input_data = tf.transpose(input_data, [1, 0, 2])
outputs = tf.TensorArray(tf.float32, input_data.shape[0])
state = self.cell.zero_state(input_data.shape[1], dtype=tf.float32)
for i in tf.range(input_data.shape[0]):
output, state = self.cell(input_data[i], state)
outputs = outputs.write(i, output)
return tf.transpose(outputs.stack(), [1, 0, 2]), state
更多有关 AutoGraph 的细节和特点,可以参考这篇指南。
tf.metrics
聚合了数据,tf.summary
记录了数据
如果要展示模型的细节,使用tf.summary.(scalar|histogram|...)
,以及使用一个情景管理器(context manager)将其重定向至一个 writer。(如果你忽略了情景管理器,并不会发生什么。)与 TensorFlow 1.x 不同,summaries 会直接送到 writer;所以不需要什么“merge”操作,也不需要调用add_summary()
,这意味着step
值必须在调用点提供。
summary_writer = tf.summary.create_file_writer('/tmp/summaries')
with summary_writer.as_default():
tf.summary.scalar('loss', 0.1, step=42)
为了把数据记录在 summaries 之前聚合它们,使用tf.metrics
。Metrics 是基于状态的:它们积累值,并且在你调用.result()
的时候返回一个渐增的结果。使用.reset_states()
来重置累积的值。
def train(model, optimizer, dataset, log_freq=10):
avg_loss = tf.keras.metrics.Mean(name='loss', dtype=tf.float32)
for images, labels in dataset:
loss = train_step(model, optimizer, images, labels)
avg_loss.update_state(loss)
if tf.equal(optimizer.iterations % log_freq, 0):
tf.summary.scalar('loss', avg_loss.result(), step=optimizer.iterations)
avg_loss.reset_states()
def test(model, test_x, test_y, step_num):
loss = loss_fn(model(test_x), test_y)
tf.summary.scalar('loss', loss, step=step_num)
train_summary_writer = tf.summary.create_file_writer('/tmp/summaries/train')
test_summary_writer = tf.summary.create_file_writer('/tmp/summaries/test')
with train_summary_writer.as_default():
train(model, optimizer, dataset)
with test_summary_writer.as_default():
test(model, test_x, test_y, optimizer.iterations)
你还可以使用 TensorBoard 来可视化 summaries:
tensorboard --logdir /tmp/summaries
使用 tf.config.experimental_run_functions_eagerly()
来调试
在 TensorFlow 2.0 中,Eager execution 可以使你一步一步地运行代码,检查数据的 shape,dtypes,values。标准的 APIs,例如tf.function
、tf.keras
,从性能和可移植性的角度考虑,被设计为使用 Graph 执行。在调试的时候,使用tf.config.experimental_run_functions_eagerly(True)
来用 Eager execution 运行代码。
例如:
@tf.function
def f(x):
if x > 0:
import pdb
pdb.set_trace()
x = x + 1
return x
tf.config.experimental_run_functions_eagerly(True)
f(tf.constant(1))
>>> f()
-> x = x + 1
(Pdb) 1
6 @tf.function
7 def f(x):
8 if x > 0:
9 import pdb
10 pdb.set_trace()
11 x = x + 1
12 return x
13
14 tf.config.experimental_run_functions_eagerly(True)
15 f(tf.constant(1))
[EOF]
在使用 Keras models 或者其他 APIs 的时候,也支持 Eager execution:
class CustomModel(tf.keras.models.Model):
@tf.function
def call(self, input_data):
if tf.reduce_mean(input_data) > 0:
return input_data
else:
import pdb
pdb.set_trace()
return input_data // 2
tf.config.experimental_run_functions_eagerly(True)
model = CustomModel()
model(tf.constant([-2, -4]))
>>> call()
-> return input_data // 2
(Pdb) l
10 if tf.reduce_mean(input_data) > 0:
11 return input_data
12 else:
13 import pdb
14 pdb.set_trace()
15 -> return input_data // 2
16
17
18 tf.config.experimental_run_functions_eagerly(True)
19 model = CustomModel()
20 model(tf.constant([-2, -4]))