2.2 理解TensorFlow 2.x

如前所述,TensorFlow 2.x推荐使用高级API,例如tf.keras,但当需要对内部细节进行更多控制时,则采用TensorFlow 1.x的典型低级API。tf.keras和TensorFlow 2.x带来一些特别的好处。让我们回顾一下。

2.2.1 即刻执行

TensorFlow 1.x定义了静态计算图。这种声明式编程可能使许多人感到困惑,因为Python本身更为动态灵活。为此,遵循Python理念,另一个流行的深度学习软件包PyTorch以更加命令式和动态化的方式定义事物:你仍然需要一个计算图,但可随时定义、更改和执行节点,而无须任何特殊会话接口或占位符。这就是所谓的即刻执行(eager execution),它意味着模型的定义是动态的,且执行是即时生效的。计算图和会话都被视为实现细节。

PyTorch和TensorFlow 2风格都继承自Chainer,Chainer是另一种“强大、灵活且直观的神经网络框架”(见https://chainer.org/)。

好消息是TensorFlow 2.x原生支持即刻执行。不再需要先静态定义一个计算图然后执行它(除非你真的想要!)。所有模型都可以动态定义并立即执行。更好的消息是,所有tf.keras的API都与即刻执行兼容。可以看出,TensorFlow 2.x在核心TensorFlow社区、PyTorch社区和Keras社区之间建立了桥梁,并充分利用了它们的优点。

2.2.2 AutoGraph

更多的好消息是TensorFlow 2.0原生支持命令式Python代码,包括if-while、print()和其他Python原生特性,并可将其转换为纯TensorFlow图代码。为什么这个功能如此有用?因为Python编码非常直观,几代程序员都习惯于这种命令式编程,但将该代码转换为运行更快且自动优化的计算图格式非常费神。而这正是AutoGraph发挥作用的时候:AutoGraph接收即刻执行式的Python代码,并将其自动转换为生成计算图的代码。因此,可以再次看出,TensorFlow 2.x在命令式、动态式和即刻执行式Python编程风格与高效的图计算之间建立了桥梁,并同时兼顾了两者的优点。

AutoGraph的使用非常容易,唯一要做的就是用特殊的装饰器tf.function来注解Python代码,如以下代码所示:

067-01

如果我们查看函数simple_nn,会发现它是用于与TensorFlow内部构件交互的特殊处理程序,而simple_function则是普通的Python处理程序:

067-02

注意,使用tf.function时,你仅需要注解某个主要函数,那么,所有其中调用的其他函数将自动且无差别地转换为一个优化过的计算图。例如,在前面的代码中,无须为linear_layer函数添加注解。可将tf.function视为即时(Just In Time, JIT)编译的标记代码。通常,你不需要看到AutoGraph自动生成的代码,但是如果你感到好奇,可使用以下代码片段查看:

067-03

该代码片段将打印以下自动生成的代码:

068-01

让我们看一个示例,该示例显示了一段相同代码在添加tf.function()装饰器注解前后的运行速度的差异。此处我们使用LSTMCell()层(详见第8章,现在将其视为运行某种深度学习的黑盒):

068-02

当执行上述代码片段时,你会看到当调用tf.function()时,运行时间会缩短一个数量级:

068-03

简而言之,tf.function注解可用来装饰Python函数和方法,从而将其转换为静态计算图的等效图,并附带所有优化操作。

2.2.3 Keras API的三种编程模型

TensorFlow 1.x提供了较低层级的API。构建模式时,首先需要创建计算图,进而编译并执行。而tf.keras提供了更高层级的API,涵盖了三种不同的编程模型:Sequential API、Functional API和Model Subclassing。创建学习模型容易得就像“将乐高积木放在一起”一样,每块“乐高积木”就是一个特定的Keras.layer。让我们看一下何时该使用Sequential、Funtional和Subclassing,并注意如何根据你的特定需求混合和匹配这三种编程模型。

1. Sequential API

Sequential API是一个非常优雅、直观且简洁的模型,适用于几乎90%的场景。在上一章中,当讨论MNIST代码时,我们介绍了使用Sequential API的示例,代码及图(见图2-2)如下所示。

069-02
069-01

图2-2 Sequential模型示例

2. Functional API

当你要构建具有更复杂(非线性)拓扑的模型时,Fun-ctional API很适用,它包含多个输入、多个输出、非顺序流的保留连接以及共享和可重用的层。每一层都是可调用的(在输入中带有张量),并且每一层都返回张量作为输出。让我们看一个示例,其中有两个单独的输入、两个单独的逻辑回归输出以及一个中间的共享模块。

现在,你无须了解各层(即乐高积木)在内部进行的操作,而只需观察非线性网络拓扑即可。还要注意,一个模块可以调用另一个模块,就像一个函数可以调用另一个函数一样:

069-03

注意,首先创建一个层,然后给它一个输入。通过tf.keras.layers.Dense(1, activation ='sig-moid', name ='prediction_a')(encoded_input_a),这两步合并在一行代码中。

非线性网络拓扑如图2-3所示。

070-01

图2-3 非线性拓扑示例

在本书的后续章节中,我们将看到多个使用Functional API的示例。

3. Model subclassing

Model subclassing提供了极大的灵活性,通常在需要自定义层时使用。换句话说,当你要构建自己的特殊乐高积木,而不是组合一些标准常用的积木时,Model subclassing很有用。在复杂性方面,它确实存在较高的成本,因此仅在真正需要时才使用。在大多数情况下,Sequential API和Functional API更合适,但是如果你希望像典型的Python / NumPy开发人员那样以面向对象的方式进行思考,则可以使用Model subclassing。

因此,为了创建自定义层,我们可以将tf.keras.layers.Layer子类化并实现以下方法:

  • __init__(可选):用于定义该层要使用的所有子层。这是构造函数,可在其中声明自定义的模型。
  • build:用于初始化层的权重。你可以使用add_weight()添加权重。
  • call:用于定义前向传递。这是自定义层被调用,并以函数式链接的地方。
  • (可选)层的序列化和反序列化可分别使用get_config()from_config()实现。

让我们看一个定制化层的示例,该层简单地将输入乘以名为kernel的矩阵(为简单起见,代码文本忽略了导入行,但在GitHub代码示例上是包含的):

071-01

一旦定制化模块MyLayer()完成定义后,就可以像任何其他模块一样对其进行组合,如下例所示,其中Sequential模型通过堆叠MyLayersoftmax激活函数进行定义:

071-02

简而言之,如果你从事构建模块的任务,则可以使用Model subclassing。

在本节中,我们已经看到tf.keras提供了更高层级的API,它具有三种不同的编程模型:Sequential API、Functional API和Model subclassing。现在让我们将注意力转移到回调,这是一个不同的功能,在使用tf.keras进行训练时非常有用。

2.2.4 回调

回调(callback)是传递给模型以扩展或修改训练期间的行为的对象。在tf.keras中经常使用到以下回调:

  • tf.keras.callbacks.ModelCheckpoint:此功能用于定期保存模型的检查点,并在出现问题时恢复。
  • tf.keras.callbacks.LearningRateScheduler:此功能用于在优化过程中动态更改学习率。
  • tf.keras.callbacks.EarlyStopping:此功能用于在运行一段时间后,当验证性能停止提高时,中断训练。
  • tf.keras.callbacks.TensorBoard:此功能用于通过TensorBoard监视模型的行为。

下例中我们使用了TensorBoard:

071-03

2.2.5 保存模型和权重

模型训练完成后,需要将权重持久化保存。可以使用以下代码片段轻松实现此目标,该代码片段将权重保存为TensorFlow的内部格式:

072-01

如果希望以Keras的格式保存(跨多后端移植),可使用:

072-02

加载权重:

072-03

除了权重之外,还可以通过以下方式将模型序列化为JSON格式:

072-04

如果愿意,也可以使用以下命令将模型序列化为YAML格式:

072-05

如果要将权重和优化参数与模型一起保存,则只需使用:

072-06

2.2.6 使用tf.data.datasets训练

使用TensorFlow 2.x的另一个好处是引入了TensorFlow数据集,它作为处理异构(大型)数据集的一种首选机制,涵盖不同类别的数据(例如音频、图像、视频、文本和翻译)。首先让我们使用pip安装tensorflow-datasets

072-07

011-02截至2019年9月,已有至少85种常用数据集,并计划在https://www.tensorflow.org/datasets/datasets中添加更多数据集。

列出所有可用的数据集,并将MNIST和元数据一起加载:

073-01

获得以下数据集列表:

073-02

另外也获得了MNIST的元信息:

073-03

有时候,从NumPy数组创建数据集也很有用。让我们看看如何在以下代码片段中使用tf.data.Dataset.from_tensor_slices()

074-01

我们还可以下载数据集,对数据进行洗牌和批处理,然后从生成器中获取一个切片,如下例所示:

074-02

注意,shuffle()是对输入数据集随机洗牌的变换,而batch()则是批量创建张量。以下代码将返回可供立即便捷使用的元组:

074-03

数据集是用于以原则性方式处理输入数据的库。操作包括:

1)创建

a. 通过from_tensor_slices(),接收单个(或多个)NumPy(或张量)并支持批处理

b. 通过from_tensors(),与上面类似,但不支持批处理

c. 通过from_generator(),从生成器函数获取输入

2)变换

a. 通过batch()将数据集按指定大小进行顺序划分

b. 通过repeat()复制数据

c. 通过shuffle(),随机洗牌数据

d. 通过map(),用一个函数对数据做映射变换

e. 通过filter(),用一个函数对数据进行过滤处理

3)迭代器

a. 通过next_batch = iterator.get_next()

数据集采用TFRecord格式,它是(任意格式)数据的一种表示方式,可轻松跨多系统移植,并独立于待训练的特定模型。简而言之,数据集相对于使用feed-dict的TensorFlow 1.0而言,灵活性更高。

2.2.7 tf.keras还是估算器

除了直接图计算和tf.keras高阶API外,TensorFlow 1.x和2.x还包括一系列称为估算器(estimator)的高阶API。通过使用估算器,你无须担心如何创建计算图或处理会话,因为估算器会帮助你以类似tf.keras的方式处理这些问题。

但估算器到底是什么呢?简而言之,它们是另一种创建或使用预制模块的方法。一个比较长的解释是,它们是针对大规模准生产环境的高效学习模型,可在单机或分布式多服务器上进行训练,并在CPU、GPU或TPU上运行,且无须重新编写模型。这些模型包括线性分类器、深度学习分类器、梯度提升树(Gradient Boosted Trees)等,将在接下来的章节中进行讨论。

让我们看一个估算器示例,该估算器用来构建含2个稠密隐藏层的分类器,每个隐藏层具有10个神经元和3个输出类:

075-01

代码feature_columns=my_feature_columns是特征列的列表,每个特征列描述了你希望模型用到的某个特征。比如,一种典型的使用方法是:

075-02

此处,tf.feature_column.numeric_column()表示实值或数字特征(https://www.tensorflow.org/api_docs/python/tf/feature_column/numeric_column)。训练高效的估算器应使用tf.Datasets作为输入。下例中对MNIST进行了加载、缩放、洗牌和批处理操作:

076-01

接下来,使用tf.estimator.train_and_evaluate()方法,并将遍历数据的input_fn作为参数传入该方法,以对估算器进行训练和评估。

076-02

TensorFlow包含用于回归和分类的估算器,这些估算器已被社区广泛采用,并将至少在整个TensorFlow 2.x的生命周期内获得支持。

然而,TensorFlow 2.x的建议是如果已经采用了它们,则应继续使用。但如果仅仅是一些程序脚本,则可使用tf.keras。截至2019年4月,估算器是完全支持分布式训练的,而对tf.keras的支持是有限的。鉴于此,可能的解决方案是使用tf.keras.estimator.model_to_estimator()tf.keras模型转换为估算器,然后再利用其对分布式训练的完全支持特性。

2.2.8 不规则张量

在继续讨论TensorFlow 2.x的优点之前,我们应该注意到TensorFlow 2.x增加了对“不规则”张量的支持,不规则张量是一种特殊的稠密张量,其大小不均匀。这在处理维度随批次变化的序列和其他数据问题(例如文本句子和分层数据)时特别有用。注意,不规则张量比填充tf.Tensor更有效,因为它不会浪费计算时间或存储空间:

077-01

2.2.9 自定义训练

TensorFlow可以为我们计算梯度(自动微分),这使得开发机器学习模型变得非常容易。如果你使用tf.keras,则可直接使用fit()训练模型,无须深入研究内部如何计算梯度的细节。但是,当你想更好地控制优化过程时,自定义训练会很有用。

目前存在多种计算梯度的方法:

1)tf.GradientTape():该类用于记录自动微分的操作。举个例子,其中用到参数persistent=True(一个布尔值,用于控制是否创建持久化梯度记录,这意味着可以多次调用该对象的gradient()方法):

077-02

2)tf.gradient_function():这将返回一个函数,该函数计算其输入函数参数在指定自变量值处的导数。

3)tf.value_and_gradients_function():这将返回输入函数的值,以及它在指定自变量处的导数列表。

4)tf.implicit_gradients():它计算输入函数的输出的梯度,结果与输出依赖的所有可训练变量有关。

让我们看一个自定义梯度计算的结构,其中,模型作为输入,训练步骤是计算total_loss=pred_loss+regularization_loss。装饰器@tf.function用于AutoGraph,tape.gradient()apply_gradients()用于计算和应用梯度:

078-01

将训练步骤train_step(inputs, labels)应用在每个epoch以及train_data中的每个输入及其关联的标签上:

078-02

简单来说,GradientTape()允许我们控制和更改内部训练过程的执行方式。在第9章中,你将看到一个更具体的使用GradientTape()训练自编码器的示例。

2.2.10 TensorFlow 2.x中的分布式训练

TensorFlow 2.x的一个非常有用的新增功能是以简单的几行额外代码,使用分布式GPU、多机器和TPU来训练模型。tf.distribute.Strategy正是用于此处的TensorFlow API,它同时支持tf.kerastf.estimator API以及即刻执行。你只需更改策略实例即可在GPU、TPU和多机器之间切换。策略可以是同步的,其中所有worker都以同步数据并行计算的方式在输入数据的不同切片上进行训练。策略也可以是异步的,其中优化器的更新不会同步进行。所有策略均要求数据由tf.data.Dataset API批量加载。

注意,分布式训练的支持功能仍处于试验阶段。路线图如图2-4所示。

078-03

图2-4 分布式训练对不同策略和API的支持

接下来详细讨论图2-4中列出的不同策略。

1. 多GPU

我们讨论了TensorFlow 2.x如何利用多个GPU。如果我们想在单台机器的多个GPU上进行同步分布式训练,则需要做两件事:©1)以一种可分发到GPU中的方式加载数据;2)将一些计算分配到GPU中:

1)以分发到GPU中的方式加载数据需要使用tf.data.Dataset(见2.2.6节)。如果没有tf.data.Dataset但有一个普通的张量,那么可用tf.data.Dataset.from_tensors_slices()将后者轻松转换为前者。这将在内存中继续使用张量,并返回源数据集,其元素是给定张量的切片。

在下面的简单示例中,我们使用NumPy生成训练数据x和标签y,然后使用tf.data.Dataset.from_tensor_slices()将其转换为tf.data.Dataset。接下来,应用洗牌以避免在跨GPU训练时产生偏差,然后生成SIZE_BATCHES批处理:

079-01

2)为了将一些计算分配到GPU中,我们实例化了一个对象distribution=tf.dis-tribute.MirroredStrategy(),该对象支持在单台机器的多GPU上执行同步分布式训练。然后,将Keras模型的创建和编译移入strategy.scope()中。注意,模型中的每个变量都在所有副本上做了镜像操作。示例如下所示:

079-02

注意,每一批次的给定输入在多GPU之间均匀分配。例如,如果将Mirrored-Strategy()策略与双GPU一起使用,那么每批大小为256的输入将在两个GPU之间分配,每个GPU每步收到128个输入样例。另外注意,每个GPU都会对收到的批次数据进行优化,而TensorFlow后台会帮助我们合并这些独立的优化结果。如果你想了解更多信息,可以在线查看notebook(https://colab.research.google.com/drive/1mf-PK0a20CkObnT0hCl9VPEje1szhHat#scrollTo=wYar3A0vBVtZ),我将解释如何在Colab(其中有针对MNIST分类而构建的Keras模型)中使用GPU。该notebook可在GitHub库中找到。

简而言之,使用多GPU非常容易,只需对用于单个服务器的tf.keras代码进行少量的更改即可。

2. 多工镜像策略

该策略在多个worker上实现同步分布式训练,每个worker都可能具有多个GPU。截至2019年9月,该策略仅适用于估算器,并且为tf.keras提供了实验性支持。如果你的目标是扩展到高性能的单台机器之外,则应使用此策略。数据必须通过tf.Dataset加载并在worker之间共享,以便每个worker都能读取唯一子集。

3. TPU策略

此策略在TPU上实施同步分布式训练。TPU是Google的专用ASIC芯片,旨在以比GPU更高效的方式显著加速机器学习负载。第16章将讨论有关TPU的更多信息。可参考此公开信息(https://github.com/tensorflow/tensorflow/issues/24412):

“要点是我们打算与TensorFlow 2.1一起宣布对TPUStrategy的支持。TensorFlow2.0将在有限用例下继续工作,但包含TensorFlow 2.1中的很多改进(bug修复、性能提升),因此我们不认为TPUStrategy已经准备好了。”

4. 参数服务器策略

该策略可实现多GPU同步本地训练或异步多机训练。对于单机本地训练,模型的变量存放在CPU上,并在所有本地GPU上复制操作。

对于多机训练,一些机器将指定为worker,还有一些机器将指定为参数服务器,并把模型变量放置在参数服务器上。计算工作在所有worker的所有GPU之间复制。可以用环境变量TF_CONFIG设置多个worker,示例如下所示:

080-01

在本节中,我们已经看到了如何使用分布式GPU、多机器和TPU以非常简单的方式以及很少的额外代码来训练模型。现在,让我们看看1.x和2.x之间的另一个区别:命名空间。

2.2.11 命名空间的改动

TensorFlow 2.x付出了巨大的努力来清理TensorFlow 1.x中极为拥挤的命名空间,尤其是根命名空间,这使得搜索变得困难。以下是主要的更改:

  • tf.keras.layers:包含以前在tf.layers下的所有符号
  • tf.keras.losses:包含以前在tf.losses下的所有符号
  • tf.keras.metrics:包含以前在tf.metrics下的所有符号
  • tf.debugging:用于调试的新命名空间
  • tf.dtypes:数据类型的新命名空间
  • tf.io:I/O的新命名空间
  • tf.quantization:用于量化的新命名空间

TensorFlow 1.x当前总共提供了2000多个端点,其中根命名空间中有500多个端点。TensorFlow 2.x删除了214个端点,其中根命名空间中删除了171个端点。TensorFlow 2.x添加了一个转换脚本,以帮助实现从1.x到2.x的转换并高亮显示已弃用的端点。

011-02新增端点和已弃用端点的完整描述可在https://github.com/tensorflow/community/blob/master/rfcs/20180827-api-names.md中找到。

2.2.12 1.x至2.x的转换

TensorFlow 1.x脚本无法直接与TensorFlow 2.x一起使用,需要进行转换。从1.x转换为2.x的第一步是使用随2.x一起安装的自动转换脚本。对于单个文件,可以使用以下命令运行它:

081-01

对于目录中的多个文件,语法为:

081-02

该脚本将尝试自动升级到2.x,并在无法升级的地方显示错误消息。

2.2.13 高效使用TensorFlow 2.x

2.x原生代码应遵循许多最佳实践:

1)默认使用tf.keras(或在某些情况下为估算器)之类的高阶API,除非有自定义操作需要,否则应避免使用直接操作计算图的低阶API。因此,通常不使用tf.Sessiontf.Session.run

2)添加装饰器tf.function,从而在基于AutoGraph的图模式下高效运行。仅用tf.function装饰高阶计算。由高阶计算调用的所有函数都将自动添加注解。如此一来,你可以做到一箭双雕:获得支持即刻执行的高阶API,以及计算图的效率。

3)使用Python对象来跟踪变量和损失。为此,使用Python语言,并使用tf.Variable而不是tf.get_variable。这样,变量将被视为普通Python作用域。

4)使用tf.data数据集进行数据输入,并将这些对象直接提供给tf.keras.Model.fit。这样,你将拥有一组用于处理数据的高性能类,并将采用最佳方法从磁盘上流式传输训练数据。

5)只要有可能,就使用tf.layers模块来组合预定义的“乐高积木”,不管是使用Sequential API、Functional API还是Subclassing编程模型。如果你需要具有准生产的模型,尤其是这些模型需要在多个GPU、CPU或多个服务器上扩展时,需使用估算器。必要时,考虑将tf.keras模型转换为估算器。

6)考虑在GPU、CPU和多个服务器上使用分布式策略。用tf.keras实现很容易。

还可以提出很多其他建议,但前述建议是最靠前的6个。TensorFlow 2.x使初始学习步骤变得非常容易,而采用tf.keras对初学者来说则非常容易。