[翻译] 更有效的 TensorFlow 2 - Effective TensorFlow 2

原  文：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]))