【干货】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_ib_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)

将为TensorBoard指定记录文件夹(tensorboard --logdir /tmp/summaries)即可将生成的记录可视化。


参考链接:

  • https://github.com/tensorflow/docs/blob/master/site/en/r2/guide/effective_tf2.md


-END-

专 · 知

专知《深度学习:算法到实战》课程全部完成!480+位同学在学习,现在报名,限时优惠!网易云课堂人工智能畅销榜首位!

欢迎微信扫一扫加入专知人工智能知识星球群,获取最新AI专业干货知识教程视频资料和与专家交流咨询!

请加专知小助手微信(扫一扫如下二维码添加),加入专知人工智能主题群,咨询《深度学习:算法到实战》课程,咨询技术商务合作~

请PC登录www.zhuanzhi.ai或者点击阅读原文,注册登录专知,获取更多AI知识资料!

点击“阅读原文”,了解报名专知《深度学习:算法到实战》课程

展开全文
Top
微信扫码咨询专知VIP会员