适合现实世界的最佳实践(加速模型训练)-平芜编程栈

加速模型训练

回想一下第7 章所述的“取得进展的循环”：想法的质量取决于这一想法经历了多少轮完善，
如图13-1 所示。你对一个想法进行迭代的速度，取决于创建实验的速度、运行实验的速度以及
分析结果数据的速度。

随着你掌握的Keras API 专业知识越来越多，深度学习实验的代码编写速度将不再是这个循
环的瓶颈。接下来的瓶颈是模型的训练速度。利用快速的训练基础设施，你可以在10 ～ 15 分
钟内得到结果。因此，你可以每天运行数十次迭代。更快的训练速度可以直接提高深度学习解
决方案的质量。
本节将介绍以下3 种提高模型训练速度的方法：

混合精度训练，利用这种方法甚至可以只用一块 GPU 来训练；
在多块 GPU 上训练；
在 TPU 上训练。
我们分别来看一下。

使用混合精度加快GPU 上的训练速度

如果我告诉你，有一种简单的技巧，可以将几乎所有模型的训练速度提高3 倍，而且基本
上是免费的，你会怎么想？这听起来过于美好而不像是真的，但这样的技巧确实存在，它就是
混合精度训练（mixed-precision training）。为了理解它的工作原理，我们首先来看一下计算机科
学中的“精度”这一概念。

1. 理解浮点数精度

精度之于数字就像分辨率之于图像。因为计算机只能处理1 和0，所以计算机中的任何数
字都需要编码为二进制字符串。比如，你可能熟悉uint8 整数，它是用8 位（bit）编码的整数：
对于uint8，00000000 表示0，11111111 表示255。要想表示大于255 的整数，需要添加更多位，
8 位是不够的。大多数整数存储用的是32 位，可以表示从−2 147 483 648 到2 147 483 647 的有
符号整数。

浮点数也是如此。在数学中，实数形成一个连续轴，任意两个数字之间都有无数个点。你
可以将实数轴不断放大。在计算机科学中，这种说法是不正确的。举个例子，在3 和4 之间的
点的个数是有限的。有多少个？这取决于你所使用的精度，即用多少位来存储一个数字。你只
能将实数轴放大到一定的分辨率。

我们通常使用以下3 种浮点数精度：

半精度或 float16，数字用16 位存储；
单精度或 float32，数字用32 位存储；
双精度或 float64，数字用64 位存储。

你可以这样来考虑浮点数分辨率：可以安全处理的任意两个数字之间的最小距离。对于单
精度，这大约是1e−7；对于双精度，这大约是1e−16；对于半精度，这则只有1e−3。

到目前为止，本书所有模型使用的都是单精度数字：模型将状态存储为float32 权重变量，
在float32 输入上运行计算。这样的精度足以在不丢失信息的情况下运行模型的前向传播和反
向传播，特别是梯度更新比较小的情况（回想一下，典型的学习率是1e−3，常见的权重更新大
小在1e−6 量级）。

你也可以使用float64，但这样做会浪费资源——像矩阵乘法或矩阵加法这样的运算在
双精度下的计算代价要大很多，所以你要做两倍的工作，却没有明显的收益。但是你不能用
float16 的权重和计算来完成同样的事情，这样的话，梯度下降过程将无法顺利运行，因为你
无法表示1e−5 或1e−6 左右的小梯度更新。

但是，你可以使用一种混合方法，即混合精度。它的思路是，在不需要太高精度的地方使
用16 位计算，而在其他地方使用32 位来保持数值稳定性。新款的GPU和TPU都具有专门的硬件，运行16 位运算比运行同等的32 位运算速度更快，占用的内存更少。通过尽可能使用这些较低
精度的运算，你可以在这些设备上大大加快训练速度。此外，在模型中对精度敏感的部分保持
使用单精度，可以在不影响模型质量的前提下获得速度收益。

这些收益是相当可观的：在新款NVIDIA GPU 上，混合精度可以将训练速度提高3 倍。在
TPU 上训练时也是有好处的（稍后会讨论这一点），可以将训练速度提高60%。

2. 混合精度训练的实践

在GPU 上训练时，混合精度的使用方法如下。

fromtensorflowimportkeras keras.mixed_precision.set_global_policy("mixed_float16")

通常情况下，模型的大部分前向传播是使用float16 完成的（除了像softmax 这样数值不
稳定的运算），而模型权重则使用float32 存储和更新。

Keras 层都有variable_dtype 和compute_dtype 这两个属性。默认情况下，这两个
属性都被设为float32。使用混合精度后，大多数层的compute_dtype 会切换为float16，
这些层会将输入转换为float16，并以float16 格式执行计算（使用半精度的权重）。然
而，这些层的variable_dtype 仍是float32，所以它们的权重能够从优化器接收到准确的
float32 更新，而不是半精度更新。

请注意，如果使用float16，那么有些运算可能会在数值上不稳定（特别是softmax 和交
叉熵）。如果你想让某一层不使用混合精度，那么只需将参数dtype=“float32” 传递给该层的
构造函数。

多GPU 训练

虽然GPU 每年都在变得越来越强大，但深度学习模型也在变得越来越大，需要越来越多的
计算资源。在单块GPU 上训练，会使训练速度受到严格限制。要解决这个问题，你可以添加更
多GPU，开始多GPU 分布式训练。

在多台设备上进行分布式计算有两种方法：数据并行（data parallelism）和模型并行（model
parallelism）。

利用数据并行，单个模型可以被复制到多台设备上。每个模型副本处理不同的数据批量，
然后将结果合并。

利用模型并行，单个模型的不同部分可以在不同设备上运行，同时处理一批数据。这种方
法对那些具有天然并行架构的模型效果最好，比如具有多个分支的模型。

在实践中，模型并行只用于那些太大而无法在单一设备上运行的模型，它不是用于加快普
通模型的训练速度，而是用于训练更大的模型。本节不会介绍模型并行，而是会重点介绍你在
大多数情况下会用到的方法：数据并行。我们来看一下它的工作原理。

1. 获得两块或多块GPU

首先，你需要能够使用多块GPU。目前，谷歌Colab 只允许使用单块GPU，所以你需要做
以下两件事之一。

获得2～4块GPU，将它们安装在一台机器上（这需要强力电源），然后安装CUDA驱动、
cuDNN 等。对大多数人来说，这并不是最佳选择。
在谷歌云、Azure 或 AWS 上租用多 GPU 虚拟机。你将能够使用带有预装驱动和软件的
虚拟机镜像，而且设置开销很小。如果你不是全天24 小时都在训练模型，那么这可能
是最佳选择。

我们不会详细介绍如何创建多GPU 云虚拟机，因为这些信息更新很快，而且很容易在网上
找到。

如果不想花费精力来管理自己的虚拟机实例，那么你还可以使用TensorFlow Cloud。这是
我和我的团队最近发布的一个软件包。你只需在Colab 笔记本的开头添加一行代码，即可开始
在多块GPU 上进行训练。如果你正在Colab 中调试模型，想无缝过渡到在尽可能多的GPU 上
训练模型，那么可以试一下这种方法。

2. 单主机、多设备同步训练

在一台具有多块GPU 的机器上，如果你能够运行import tensorflow，那么很快就可以
开始训练分布式模型。它的工作原理如下。

这几行代码实现了最常见的训练设置：单主机、多设备同步训练，在TensorFlow 中也叫“镜
像分布式策略”。“单主机”是指所有GPU 都在一台机器上（与之相对的是由许多机器组成的集
群，每台机器都有自己的GPU，机器之间通过网络进行通信）。“同步训练”是指所有GPU 模
型副本的状态始终保持相同——有些分布式训练的变体并非如此。

开启MirroredStrategy 作用域并在其中构建模型时，MirroredStrategy 对象会在每
块可用的GPU 上创建一个模型副本。然后，每个训练步骤都以如下方式进行（参见图13-2）

(1) 从数据集中抽取一批数据（叫作全局批量）。
(2) 将这批数据分为4 个子批量（叫作局部批量）。举个例子，如果全局批量包含512 个样本，
那么每个局部批量都有128 个样本。我们希望局部批量足够大，能够保持GPU 持续运转，
所以全局批量一般都很大。
(3) 4 个副本中的每一个都在自己的设备上独立处理一个局部批量，即运行一次前向传播和
一次反向传播。每个副本输出一个“权重增量”，其含义是：给定模型在局部批量上的
损失相对于权重的梯度，模型每个权重变量的更新大小。
(4) 将4 个副本得到的局部梯度权重增量合并，得到一个全局增量，并将其应用于所有副本。
由于这是在每一步结束时进行的，因此各个副本总是保持同步，它们的权重始终保持
相同。

在理想情况下，在N 块GPU 上进行训练会带来N 倍的速度提升。然而在实践中，分布式
会引入一些开销，特别是合并来自不同设备的权重增量需要一些时间。你获得的有效加速取决
于所使用的GPU 数量：

使用 2 块 GPU，速度约为 2 倍；
使用 4 块 GPU，速度约为 3.8 倍；
使用 8 块 GPU，速度约为 7.3 倍。
这里假设全局批量足够大，以保持每块GPU 满负荷运转。如果批量太小，那么局部批量将
不足以保持GPU 持续运转。

TPU 训练

除了使用GPU，深度学习领域还有一种趋势，那就是将工作流程转移至日益专业化的硬
件上。这些硬件都是专门为深度学习工作流程设计的。这种单一用途的芯片叫作专用集成电路
（application-specific integrated circuit，ASIC）。许多家大大小小的公司正在开发新的芯片，但目
前在这方面最重要的成果是谷歌的张量处理单元（Tensor Processing Unit，TPU），它可以在谷歌
云和谷歌Colab 上使用。

在TPU 上训练需要跨越一些障碍，但这些额外的工作是值得做的，因为TPU 真的非常快。
在TPU V2 上训练通常要比在NVIDIA P100 GPU 上快15 倍。对于大多数模型而言，TPU 训练
的成本效益平均比GPU 高3 倍。

1. 通过谷歌Colab 使用TPU

你可以在Colab 中免费使用8 核TPU。在Colab 菜单的Runtime（代码执行程序）标签下，
单击Change Runtime Type（更改运行时类型），你会发现除了GPU 运行时，还可以选择TPU 运
行时。

使用GPU 运行时，模型可以直接访问GPU，无须任何特殊操作。如果使用TPU 运行时，
则并非如此。在开始构建模型之前，你还需要进行额外操作：连接到TPU 集群。

它的具体代码如下。

importtensorflowastf tpu=tf.distribute.cluster_resolver.TPUClusterResolver.connect()print("Device:",tpu.master())

你不必过分担心这段代码的作用——它只不过是一句“芝麻开门”，将笔记本运行时连接到
设备。

与多GPU 训练的情况类似，使用TPU 训练也需要开启分布式策略作用域，即TPUStrategy
作用域。TPUStrategy 遵循与MirroredStrategy 相同的分布式模板，模型在每个TPU 内核
上复制一次，并且各个副本保持同步。

下面来看一个简单的例子，如代码清单13-4 所示。
代码清单13-4 在TPUStrategy 作用域中构建模型

fromtensorflowimportkerasfromtensorflow.kerasimportlayers strategy=tf.distribute.TPUStrategy(tpu)print(f"Number of replicas:{strategy.num_replicas_in_sync}")defbuild_model(input_size):inputs=keras.Input((input_size,input_size,3))x=keras.applications.resnet.preprocess_input(inputs)x=keras.applications.resnet.ResNet50(weights=None,include_top=False,pooling="max")(x)outputs=layers.Dense(10,activation="softmax")(x)model=keras.Model(inputs,outputs)model.compile(optimizer="rmsprop",loss="sparse_categorical_crossentropy",metrics=["accuracy"])returnmodelwithstrategy.scope():model=build_model(input_size=32)

我们几乎可以开始训练了。但是Colab 中的TPU 有些古怪：它是双虚拟机设置，也就是说，
托管笔记本运行时的虚拟机与TPU 所在的虚拟机不同。因此，你无法利用本地磁盘（与托管笔
记本的虚拟机相连的磁盘）存储的文件进行训练。TPU 运行时无法从本地磁盘进行读取。加载
数据的方法有以下两种。

利用虚拟机内存（而不是磁盘）中的数据进行训练。如果数据是 NumPy 数组，那么你
已经完成这一步了。
将数据存储在谷歌云存储（Google Cloud Storage，GCS）的存储桶中，并创建一个数据
集，直接从存储桶中读取数据，无须下载到本地。TPU 运行时可以从GCS 中读取数据。
大型数据集无法整个存储于内存中，只能选择这种方法。

在本例中，我们利用内存中的NumPy 数组（CIFAR10 数据集）进行训练。

你会发现，需要一段时间才会开始第一轮训练，这是因为需要先将模型编译为TPU 可处理
的对象。完成这一步之后，训练速度就会变得飞快。

2. 利用步骤融合来提高TPU 利用率

TPU 拥有强大的计算能力。应该使用非常大的批量来训练，以保持TPU 内核持续运转。对
于小模型来说，批量可能会变得非常大，每个批量超过10 000 个样本。在处理特别大的批量时，
你应该相应地增大优化器的学习率。你还应该减少权重的更新次数，但每次更新都会更加准确（因
为梯度是用更多数据点计算得到的），所以每次更新时权重的变化幅度应该更大。

然而，你可以用一个简单的技巧来保持合理的批量大小，同时充分使用TPU，它就是步骤
融合（step fusing）。具体做法是在每个TPU 执行步骤中运行多个训练步骤，也就是说，在从
虚拟机内存到TPU 的两次往返之间做更多工作。要实现这一技巧，只需在compile() 中指定
steps_per_execution 参数，比如steps_per_execution=8 就是在每个TPU 执行步骤中
运行8 个训练步骤。对于没有充分利用TPU 的小模型来说，这可以带来巨大的速度提升。