【干货】TensorFlow 2.0官方风格与设计模式指南（附示例代码）- 专知

【干货】TensorFlow 2.0官方风格与设计模式指南（附示例代码）

【导读】TensorFlow 1.0并不友好的静态图开发体验使得众多开发者望而却步，而TensorFlow 2.0解决了这个问题。不仅仅是默认开启动态图模式，还引入了大量提升编程体验的新特性。本文通过官方2.0的风格指南来介绍新版本的开发体验。

TensorFlow 2.0做了大量的改进来提升开发者的生产力，移除了冗余的API，让API更加一致（统一的RNN、统一的优化器），将动态图模式（Eager Execution）与Python运行时集成地更加紧密。

下面先简单介绍一下主要的变更：

API清理

TensorFlow 2.0删除或移动了许多API。例如，删除了tf.app、tf.flags和tf.logging，将tf.contrib下的工程搬家。通过将低频使用的方法放到子包的方法来清理tf.*，例如tf.math。一些API被替换成了等价API，如tf.summary、tf.keras.metrics和tf.keras.optimizers。

Eager Execution（动态图模式）

TensorFlow 1.X 要求用户手动构建静态图，并通过sess.run来执行。而TensorFlow 2.0可以像Python普通程序那样直接执行，其中的Graph和Session更像是实现细节。

Eager模式使得tf.control_dependencies()不再被需要，因为代码会按照代码顺序执行。（使用tf.function时，有副作用的代码会按照代码顺序执行）。

TensorFlow 1.X 要求用户手动构建静态图，并通过sess.run来执行。而TensorFlow 2.0可以像Python普通程序那样直接执行，其中的Graph和Session更像是实现细节。

不再有全局

TensorFlow 1.X 非常依赖于隐式的全局命名空间，当你调用tf.Variable()时，变量会被放到默认图中，就算你丢失了指向它的Python变量，它依然会存在。之后你可以通过通过它的变量名来恢复它。当你并不能控制变量的创建时，这就变得非常艰难。因此，许多机制都在帮助用户找回变量和帮助框架找回用户创建的变量：Variable scopes、global collections、一些帮助函数如tf.get_global_step()、tf.global_variables_initializer()以及优化器也在隐式地为所有可训练变量计算梯度等。TensorFlow 2.0删除了所有这些机制，而采用了默认机制：跟踪你自己的变量！如果你丢失了对某个变量的跟踪，它会被垃圾回收机制回收。

这样的机制给用户增加了额外的工作，但使用Keras对象会减轻用户的负担。

函数，不是会话

调用session.run()几乎像是一个函数调用：你指定输入和需要调用的函数，然后你得到输出集合。在TensorFlow 2.0中，你可以用tf.function来装饰一个Python函数来使用JIT编译，这样TensorFlow会将它当成一个单独的图来执行。这使得TensorFlow可以得益于图模式：

性能：函数可以被优化（节点剪枝、核融合等）
便携式：函数可以被导出/导入，用户可以复用和分享模块化的TensorFlow函数

# TensorFlow 1.X
outputs = session.run(f(placeholder), feed_dict={placeholder: input})
# TensorFlow 2.0
outputs = f(input)

由于用户可以将Python和TensorFlow代码混写，我们希望用户可以充分利用Python的表达性。但是便携式的TensorFlow要在没用Python解释器的环境下运行 - 移动端、C++和JS。为了避免用户重写代码，当使用@tf.function时，AutoGraph会将Python结构的子集转换为TensorFlow等价物：

for/while -> tf.while_loop （支持break和continue）
if -> tf.cond
for _ in dataset -> dataset.reduce

AutoGraph支持嵌套的控制流，使得许多复杂机器学习的开发变得精简，且能保证效率，例如序列模型、强化学习、定制化的训练循环等。

下面介绍TensorFlow 2.0的风格和设计模式：

将代码重构为一些小函数

TensorFlow 1.X的常见用例模式是"kitchen sink"策略，所有可能的计算都被事先统一构建好，然后用session.run()来执行所选的张量。在TensorFlow 2.0中，用户应该讲代码按需重构为一些小函数。一般情况下，并不需要将所有小函数用tf.function来装饰；只要用tf.function来装饰高级计算 - 例如训练的一步、或者模型的前向传播。

用Keras层和模型来管理变量

Keras模型和层提供了便利的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)
  ...

# 你仍需要管理w_i和b_i，并且它们的形状的定义在代码的其他地方K

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层和模型都继承自tf.train.Checkpointable并且与@tf.function集成，使得用Keras对象直接保存和导出SavedModel变得可能。你并不需要使用Keras的fit() API来使用这些集成特性。

这里有一个迁移学习的例子，可以展现Keras如何轻松地收集相关变量子集。比如你正在训练一个共享主干的multi-headed的模型：

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.gradients(loss, path1.trainable_variables)
  optimizer.apply_gradients(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.gradients(loss, head2.trainable_variables)
  optimizer.apply_gradients(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是最好的从硬盘流式使用训练数据的方法。Datasets是iterables （不是iterators），和Eager模式下其他Python的iterables类似。通过tf.function()来封装你的代码，可以充分利用数据集异步预抓取/流式特性，它会用AutoGraph将Python迭代器替换为等价的图操作。

@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.gradients(loss, model.trainable_variables)
    optimizer.apply_gradients(gradients, model.trainable_variables)

如果用的是Keras的.fit() API，你不必关心数据集迭代：

model.compile(optimizer=optimizer, loss=loss_fn)
model.fit(dataset)

利用AutoGraph和Python控制流

AutoGraph提供了一种将依赖数据的控制流转换为图模式的等价物，如tf.cond和tf.while_loop。

序列模型中经常出现依赖数据的控制流。tf.keras.layers.RNN封装了RNN单元，让你可以静态或动态地来展开循环。你可以将动态展开实现如下：

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的特性可以在下面链接中查看：

https://github.com/tensorflow/docs/blob/master/site/en/r2/guide/autograph.ipynb

tf.metrics来合计数据和用tf.summary来记录数据

完整的tf.summary符号即将推出。你可以通过下面方法来使用TensorFlow 2.0的tf.summary：

from tensorflow.python.ops import summary_ops_v2

你可以使用tf.summary.(scalar|histogram|...)来记录数据，独立使用它时并不会做任何事情，你需要利用上下文管理器将它重定向到合适的file writer。（这避免了硬编码将日志写入特定文件）

summary_writer = tf.summary.create_file_writer('/tmp/summaries')
with summary_writer.as_default():
  summary_ops_v2.scalar('loss', 0.1, step=42)

为了在记录前合计数据，你可以使用tf.metrics。Metrics是有状态的，它们会累积值并在你调用.reuslt()方法时返回一个累计结果。你可以用.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):
      summary_ops_v2.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)
  summary_ops_v2.scalar('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)