基于 GCP 的分布式训练 TensorFlow 模型入门指南
文 / Nikita Namjoshi，机器学习解决方案工程师
对于数据科学领域的许多人而言，分布式训练似乎是一项令人望而生畏的任务。除了构建和仔细评估高质量机器学习模型外，您还必须了解如何针对特定硬件优化模型以及管理基础架构。数据科学家的工具包中通常不包含后一种技能。但是，借助 Google Cloud Platform (GCP) 上的代管式服务，您可以将模型训练作业轻松扩展到多个加速器甚至多台机器，而无需 GPU 专业知识。
本文以教程风格编写，旨在帮助您获得 GCP 数据科学工具的实践经验，并跨多个 GPU 训练 TensorFlow 模型。此外，您还将了解一些分布式训练领域的关键术语，例如数据并行、同步训练和 AllReduce。
数据并行是您将在本文中学习的概念之一
为什么要使用分布式训练？
每个数据科学家和机器学习工程师都有过坐着痛苦等待模型训练完成的经历。即使您有权使用 GPU，但对于大型数据集，我们也可能需要数天的时间来等待大型深度学习模型收敛。使用正确的硬件配置可以将训练时间缩短至数小时甚至数分钟。更短的训练时间有助于加快迭代速度，从而实现您的建模目标。
如果您有可用的 GPU，TensorFlow 会自动使用它，无需更改任何代码。类似地，TensorFlow 可以不经设置便利用多个 CPU 核心。但是，如果想要使用两个或多个 GPU 进行训练，那么必须完成一些额外的工作。这项额外工作必不可少，因为 TensorFlow 需要了解如何在运行时协调跨 GPU 训练过程。幸运的是，借助 tf.distribute 模块，您可以获取不同的分布式训练策略，而且您可以将这些策略轻松融合到程序中。
在进行分布式训练时，明确机器与设备之间的区别至关重要。设备是指 TensorFlow 可以在上面运行运算的某台机器上的 CPU 或加速器（例如 GPU 或 TPU）。本文重点介绍如何使用一台拥有多个 GPU 设备的机器进行训练，不过，tf.distribute.Strategy API 也提供了对多进程训练的支持。在多进程设置中，会将训练分布到多台机器上。这些机器可以仅包含 CPU，也可以每台机器有一个或多个 GPU 设备。
单 GPU 训练
在下面的 Colab Notebook 中，您将找到基于 Cassava 数据集 训练 ResNet50 架构 的代码。如果您执行 Colab 中的单元并训练模型，您会注意到，每个周期执行的步骤数为 89，每个周期约要 100 秒。请记下这些数字；我们稍后再来讨论。
多 GPU 训练
您可以在 Colab 中访问单个 GPU，但如果您想使用多个 GPU，那就无能为力了。此外，尽管 Colab Notebook 非常适合快速实验，但您可能需要更安全、更可靠的设置，以便更好地控制环境。为此，您可以转向云上环境 GCP。
您可以利用许多不同的方法在 GCP 上执行分布式训练。如果您是进行实验的学生/研究员，而不是在生产工作流中训练模型的公司工程师，那么为您的用例选择最佳选项可能会涉及不同的考虑因素。
在本文中，您将使用 GCP AI Platform Notebooks。这种途径不但为分布式训练提供了一种简单方法，还让您有机会探索在 GCP 上运行的代管式笔记本环境。或者，如果您已经设置了本地环境，并且希望在本地和 GCP 环境之间实现自由转换，那么请了解一下 TensorFlow Cloud 库。TensorFlow Cloud 可以自动执行本文介绍的许多步骤；但是，我们将在这里逐一介绍步骤，以便您可以更深入地了解分布式训练中涉及的关键概念。
在接下来的部分中，您将学习如何使用 tf.distribute.Strategy API 修改单 GPU 训练代码。生成的代码将与云平台无关，因此您可以在其他环境中运行这些代码而无需任何更改。另外，您也可以在自己的硬件上运行相同的代码。
为分布式训练准备代码
使用 tf.distribute.Strategy API 的第一步是实例化您的策略。在本教程中，您将使用 MirroredStrategy，这是 TensorFlow 中提供的几种分布策略之一。
strategy = tf.distribute.MirroredStrategy()
接下来，您需要在策略的范围内封装模型参数的创建过程。这一步至关重要，因为它会告知 MirroredStrategy 在您的 GPU 设备上镜像复制哪些变量。
with strategy.scope():
model = create_model()
model.compile(
loss='sparse_categorical_crossentropy',
optimizer=tf.keras.optimizers.Adam(0.0001),
metrics=['accuracy'])
在运行更新的代码之前，我们简单了解一下在我们调用 model.fit 时实际发生的情况，以及由于添加了策略而导致训练有哪些不同。为简单起见，假设您有一个简单的线性模型而不是 ResNet50 架构。在 TensorFlow 中，您可以从其计算图的角度来考虑这个简单的模型。
在下图中，您可以看到 matmul 运算获取 X 和 W 张量，它们分别是训练批次和权重。随后，生成的张量与张量 b（模型的偏差项）一起传递给加法运算。此运算的结果为 Y pred，即模型的预测值。
我们需要一种方法来执行此计算图，以便可以利用两个 GPU。您可以通过多种不同的方法来实现这一目标。例如，您可以将模型的不同层置于不同的机器或设备上，这是模型并行的一种形式。另外，您也可以分发数据集，以便每个设备在每个训练步骤中使用相同的模型处理输入批次的一部分，这称为数据并行。或者，您可以将两者结合使用。数据并行是最常见（也是最简单）的方法，也是我们要在这里采用的方法。
下图显示了一个数据并行的示例。输入批次 X 被一分为二，一个切片发送到 GPU 0，另一个切片发送到 GPU 1。在这种情况下，每个 GPU 都计算相同的运算，但使用数据的不同切片。
MirroredStrategy 是一种数据并行策略。因此，当我们调用 model.fit 时，MirroredStrategy 会将 ResNet50 模型复制（称为副本）到两个 GPU 上。CPU（主机）负责准备 tf.data.Dataset 批次并将数据发送到 GPU（设备）。
随后的梯度更新将以同步方式发生。这意味着每个工作进程设备都会基于输入数据的不同切片计算通过模型的前向和后向传递。随后，将这些切片中每个切片的计算梯度跨所有设备聚合，并以称为 AllReduce 的过程进行汇总（通常是平均值）。然后，优化器使用这些汇总梯度执行参数更新，从而使设备保持同步。由于每个工作进程只有在所有其他工作进程都完成当前步骤之后才能继续下一个训练步骤，因此这种梯度计算成为同步策略的分布式训练中的主要开销。
虽然 MirroredStrategy 是一种同步策略，但数据并行策略也可以是异步的。在异步数据并行策略中，每个工作进程根据输入数据的一个切片计算梯度，并以异步方式对参数进行更新。与同步策略相比，异步训练具有容错优势，因为工作进程之间并不相互依赖，但是可能导致过期的梯度。您可以通过尝试 TensorFlow 参数服务器策略 来详细了解异步训练。
通过实例化 MirroredStrategy 的这两个简单的步骤，然后在策略范围内封装模型的创建过程，TensorFlow 将通过数据并行和同步梯度更新完成跨 GPU 分配训练作业这一繁重任务。
您最后要更改的是批次大小。
BATCH_SIZE = 64 * strategy.num_replicas_in_sync
回想一下，在单 GPU 的情况下，批次大小为 64。这意味着在模型训练的每个步骤中，处理了 64 个图像，并且每个周期中生成的步骤数等于总数据集大小/批次大小，我们之前提到此数字为 89。
当您使用 tf.distribute.Strategy API 和 tf.data 进行分布式训练时，批次大小现在是指全局批次大小。换句话说，如果传递一个大小为 10 的批次，而您有两个 GPU，则每台机器每步将处理 5 个样本。在这种情况下，10 称为全局批次大小，而 5 称为 每个 副本批次大小。为了充分利用您的 GPU，您将需要根据副本的数量来缩放批次大小，在这种情况下，副本数为两个，因为每个 GPU 上都有一个副本。
您可以自己更改这些代码，也可以只使用已经进行了更改的 其他 Colab Notebook。尽管 MirroredStrategy 专为多 GPU 环境而设计，但您实际上可以在 GPU 运行时或 CPU 运行时的 Colab 中运行此笔记本而不出现错误。TensorFlow 可以不经设置便使用单个 GPU 或多个 CPU 核心，因此您实际上并不需要策略，但这对于测试/实验目的而言可能十分方便。
设置 GCP 项目
现在我们已经进行了必要的代码更改，下一步是设置 GCP 环境。为此，您需要一个已启用结算功能的 GCP 项目。
接下来，您应当 启用 Cloud Compute Engine API。如果您正在处理一个全新的项目，那么此过程可能还会提示您连接您创建的结算帐号。如果您使用的是已经处理过的 GCP 项目，则很可能已经启用了 Compute Engine API。
申请配额
Google Cloud 对资源使用量实施配额制，以防止滥用和意外使用。如果您需要访问的特定资源超过默认提供的资源，则必须申请更多配额。在本教程中，我们将使用 NVIDIA T4 Tensor Core GPU。默认情况下，您可以在每个位置获取一个 NVIDIA T4 Tensor Core GPU，但是，要进行分布式训练，您需要为该位置的其他 GPU 申请配额。
在 GCP 控制台中，滚动到左侧的汉堡式菜单，然后导航至 IAM & Admin > Quotas
在 Quotas 页面上，您可以为 Compute Engine API 添加服务过滤器。请注意，如果您尚未启用 Compute Engine API 或未启用结算功能，则不会将 Compute Engine API 视为过滤器选项，因此请确保已首先完成前述步骤。
在列表中找到 NVIDIA T4 GPU 资源后，继续并点击该行的 ALL QUOTAS。
进入 NVIDIA T4 GPU 的 Quota metric details 页面后，选择“Location: us-west1”，然后点击页面顶部的 EDIT QUOTAS。
如果您已经拥有用于不同类型 GPU 或位于其他位置的配额，则可以轻松使用这些配额。只需确保记住 GPU 的类型和位置，因为以后在设置 AI Platform Notebook 环境时需要指定这些参数。此外，如果您倾向于循序渐进，只使用一个 GPU 而不是申请两个 GPU 的配额，也可以这样做。您的代码不会被分发，但是您仍将从学习如何设置 GCP 环境中受益。
在 Quota changes 菜单中填写您的详细联系信息，然后将 New Limit 设置为 2。完成后，点击 Done。
提交申请后，您首先会收到一封确认电子邮件，随后当申请获得批准时，您会再次收到一封确认电子邮件。
创建 AI Platform Notebook 实例
在等待配额批准时，下一步是使用 AI Platform Notebooks 进行设置，可以使用与之前控制台中相同的汉堡式菜单找到该工具，然后滚动到 Artificial Intelligence > AI Platform > Notebooks
如果您是第一次使用该工具，则需要启用 API。
AI Platform Notebooks 是一种用于执行数据科学工作的代管式服务。如果您喜欢在笔记本环境中进行开发，则此工具是一个理想选择。您可以轻松添加和移除 GPU，而不必担心 GPU 驱动程序的安装，而且根据您的用例，会提供许多实例映像供您选择，因此您无需费力设置完成工作所需的所有 Python 软件包。
启用 Notebooks API 后，下一步是创建实例。您可以通过点击页面顶部的 NEW INSTANCE 按钮，然后选择 TensorFlow Enterprise 2.3 映像（如果您日后持续关注，则选择最新的 TensorFlow 映像），随后选择 With 1 NVIDIA Tesla T4。TensorFlow 企业版 是针对 GCP 优化的 TensorFlow 发行版。
点击 New notebook instance 窗口底部的 ADVANCED OPTIONS，然后更改以下字段：
- Instance name：为您的实例命名
- Region：us-west1
- GPU type：NVIDIA Tesla T4
- Number of GPUs：2
- 选中 I nstall NVIDIA GPU driver automatically for me 复选框
然后点击 CREATE。请注意，如果尚未获得 NVIDIA T4 GPU 配额的批准，点击 CREATE 时会收到一条错误消息。因此，请确保您已收到批准消息，然后再完成此步骤。此外，如果您打算在 us-west1 中使用除 NVIDIA T4 以外的其他 GPU 类型或位置，则在创建笔记本时需要更改这些参数。
您的实例将需要几分钟来启动，完成后，您会看到以蓝色字母显示的 OPEN JUPYTERLAB 选项。
请注意，即使在创建 AI Platform Notebook 实例后，也可以更改硬件（例如，添加或移除 GPU）。如果将来需要这样做，只需停止实例并 按照此处的步骤操作。
在 AI Platform Notebooks 上训练多 GPU 模型
现在，您的实例已设置完成，您可以点击 OPEN JUPYTERLAB。
以 .ipynb 文件的形式下载 Colab 笔记本，然后将其上传到 Jupyter Lab 环境。上传文件后，转到笔记本并运行代码。
执行 model.fit 单元时，您应当注意到每个周期的步骤数现在为 45，该数值为使用单个 GPU 时的一半。这表示数据并行在起作用。全局批次大小为 64 * 2，您的 CPU 将向每个 GPU 发送 64 个图像的批次。因此，虽然该模型以前在一个步骤中只能看到 64 个样本，但现在在每个步骤中都能看到 128 个样本，因此每个周期需要的时间更少。以前，每个周期大约需要 100 秒，而现在每个周期大约需要 60 秒。您会注意到，添加第二个 GPU 未能将时间减半，因为同步梯度会涉及一些开销。使用更大的数据集（木薯数据集只有 5656 个训练图像）时，好处会变得更加明显。此外，您可以使用许多技术来从第二个 GPU 中获得更多益处，例如确保您的输入流水线不会成为瓶颈。要详细了解如何充分利用 GPU，请参阅 TensorFlow 性能调试指南。
DLVM 上长时间运行的作业
到目前为止，您已经学习了如何使用 GCP AI Platform Notebooks 来执行简单的分布式训练作业。我们使用的数据集不是很大，并且仅经过了几个周期，模型便达到相当高的准确率。但是，实际上，您的训练作业可能会运行更长的时间，并且您可能不想使用笔记本。
启动 AI Platform Notebook 时，它将使用 GCP Deep Learning VM Image 创建 Google Compute Engine (GCE) 实例。Deep Learning VM Image 是针对数据科学和机器学习任务优化的 Compute Engine 虚拟机映像。在我们的示例中，我们使用了 TensorFlow Enterprise 2.3 映像，但是也 提供了许多其他选项。
在控制台中，您可以使用菜单导航到 Compute Engine > VM instances
此时，您应当看到一个实例，其名称与您之前创建的笔记本相同。由于这是一个 GCE 实例，因此我们可以 ssh 到机器并在其中运行代码。
安装 Google SDK
安装 Google Cloud SDK 后，您可以从终端管理项目中的 GCE 资源。请按照 此处的步骤 安装SDK 并连接到您的项目。
SSH 到虚拟机
安装并配置 SDK 后，您可以在终端中使用以下命令 ssh 到您的虚拟机。只需确保更改实例名称和项目名称即可。
gcloud compute ssh {your-vm-name} --project={your-project-name}
如果您在虚拟机上运行命令 nvidia-smi，您将看到我们之前配置的两个 NVIDIA T4 Tensor Core GPU。
要运行分布式训练作业，只需从 Colab 笔记本中以 .py 文件形式下载代码，然后从本地机器使用以下命令将其复制到您的虚拟机中。
gcloud compute scp --project {your-project-name} {local-path-to-py-file} {your-vm-name}:~/
最后，您可以使用以下命令在虚拟机上运行脚本
python dist_strat_blog_multi_gpu.py
此时，您应当看到模型训练作业的输出。
如果您选择笔记本环境，则可以继续使用上一部分中使用的工作流。但是，如果您更喜欢使用 vim 或 emacs，或者要使用 Screen 运行长时间运行的作业，则可以选择从终端 ssh 到虚拟机。请注意，您也可以 从命令行直接 启动深度学习虚拟机，而不是像本教程中那样使用 AI Platform Notebooks UI。
完成实验后，别忘了关闭实例。为此，您可以从 Notebook instances 页面或者控制台界面的 GCE Instances 页面中选择实例，然后点击窗口顶部的 STOP。关闭实例至关重要，因为实例每运行一个小时，您便需要相应支付数美元。您可以轻松停止实例，然后在要运行更多实验时将其重新启动，此时所有文件依然存在。
让您的分布式训练技能更上一层楼
在本文中，您学习了如何使用 MirroredStrategy（一种同步数据并行策略）在 GCP 上的两个 GPU 之间分配 TensorFlow 训练作业。现在，您已经了解如何设置 GCP 环境和准备代码的基本流程，但在分布式训练领域中，还有许多内容亟待探索。例如，如果您有兴趣将分布式训练作业构建到生产机器学习流水线中，请了解 AI Platform Training 服务，这项服务还允许您跨多台机器配置训练作业，其中每台机器都包含多个 GPU。
在 tensorflow.google.cn 网站上，您可以在 概览指南 中查看适用于 tf.distribute.Strategy API 的其他策略，还可以学习如何将策略与 自定义训练循环 结合使用。
有关更多高级概念，请参阅 关于如何分布数据的指南，以及关于 如何使用 TensorFlow Profiler 进行性能调试的指南，以确保您最大程度地利用 GPU。
原文：Getting Started with Distributed TensorFlow on GCP
中文：TensorFlow 公众号