作者：Nikita Namjoshi，机器学习解决方案工程师
对于许多数据科学领域的人来说，分布式训练似乎是一项艰巨的任务。除了构建和周到地评估高质量的 ML 模型之外，您还必须了解如何针对特定硬件优化模型并管理基础设施。后者的技能通常不包含在数据科学家的工具箱中。但是，借助 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 设备。
在 以下 Colab 笔记本 中，您将找到训练 ResNet50 架构 的代码在 木薯数据集 上。如果您在笔记本中执行单元格并训练模型，您会注意到每个纪元执行的步骤数为 89，每个纪元大约需要 100 秒。请注意这些数字；我们将在后面讨论它们。
您可以在 colab 中访问单个 GPU，但如果您想使用多个 GPU，您的运气就到此为止了。此外，虽然 Colab 笔记本非常适合快速实验，但您可能希望获得更安全、更可靠的设置，让您能够更好地控制您的环境。为此，您可以转向云。
在 GCP 上进行分布式训练有很多不同的方法。选择最适合您的用例的选项可能会涉及不同的考虑因素，例如，您是运行实验的学生/研究人员，还是一家公司的工程师在生产工作流程中训练模型。
在本文中，您将使用 GCP AI Platform 笔记本。这条路径提供了一种简单的方法来进行分布式训练，并让您有机会探索在 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 op 接收X 和W 张量，它们分别是训练批次和权重。然后将生成的张量与张量b（即模型的偏差项）一起传递给 add op。此 op 的结果是Ypred，即模型的预测。
我们想要一种执行此计算图的方法，以便我们可以利用两个 GPU。我们可以用多种不同的方法实现这一点。例如，您可以将模型的不同层放在不同的机器或设备上，这是一种模型并行。或者，您可以将数据集分布起来，以便每个设备在每个训练步骤上处理输入批次的一部分，使用相同的模型，这被称为数据并行。或者，您可能同时进行两种操作。数据并行是最常见（也是最容易）的方法，这就是我们将在这里使用的方法。
下一张图片显示了数据并行的示例。输入批次 X 被分成两半，一半发送到 GPU 0，另一半发送到 GPU 1。在这种情况下，每个 GPU 计算相同的操作，但处理数据不同的切片。
MirroredStrategy 是一种数据并行策略。因此，当我们调用 model.fit 时，MirroredStrategy 将在两个 GPU 上创建 ResNet50 模型的副本（称为副本）。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，您需要将批次大小乘以副本数量，在本例中为 2，因为每个 GPU 上都有一个副本。
您可以自己进行这些代码更改，或者简单地使用此 其他 Colab 笔记本，其中已经进行了更改。虽然 MirroredStrategy 是为多 GPU 环境设计的，但您实际上可以在 Colab 中的 GPU 运行时或 CPU 运行时上运行此笔记本，而不会出错。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 和管理 > 配额
在配额页面上，您可以添加 Compute Engine API 的服务过滤器。请注意，如果您尚未启用 Compute Engine API 或计费，则您将不会看到 Compute Engine API 作为过滤器选项，因此请确保您已先完成前面的步骤。
当您在列表中找到 NVIDIA T4 GPU 资源时，请继续点击该行的所有配额。
进入 NVIDIA T4 GPU 的配额指标详细信息页面后，选择位置：us-west1，然后点击页面顶部的编辑配额。
如果您已经拥有不同类型 GPU 或不同位置的配额，则可以轻松使用它们。只需确保记住 GPU 类型和位置，因为您将在稍后设置 AI Platform Notebook 环境时需要指定这些参数。此外，如果您希望继续并只使用单个 GPU 而不是申请两个 GPU 的配额，您也可以这样做。您的代码不会是分布式的，但您仍然可以从学习如何设置 GCP 环境中获益。
>
在配额更改菜单中填写您的联系信息，然后将新的配额设置为 2。完成后点击完成。
提交申请后，您将首先收到一封确认电子邮件，然后在您的申请获得批准后，您将收到另一封电子邮件。
创建 AI Platform Notebook 实例在等待配额审批的同时，下一步是使用 AI Platform Notebooks 进行设置，您可以在控制台的同一个汉堡包菜单中找到它，然后滚动到人工智能 > AI Platform > Notebooks
如果您是第一次使用该工具，则需要启用 API。
AI Platform Notebooks 是一种用于进行数据科学工作的托管服务。如果您喜欢在笔记本环境中进行开发，那么此工具非常理想。您可以轻松地添加和删除 GPU，而无需担心 GPU 驱动程序的安装，并且您可以根据您的用例选择多种实例映像，因此您无需费心设置完成工作所需的所有 Python 包。
启用 Notebooks API 后，下一步是创建您的实例。您可以通过点击页面顶部的“新建实例”按钮，然后选择 TensorFlow Enterprise 2.3 映像（如果您在稍后的日期继续使用，则选择最新的 TensorFlow 映像），并选择 1 NVIDIA Tesla T4 选项。 TensorFlow Enterprise 是针对 GCP 优化的 TensorFlow 分发版。
点击新建笔记本实例窗口底部的“高级选项”，然后更改以下字段
然后点击创建。请注意，如果您尚未获得 NVIDIA T4 GPU 配额的批准，则在点击“创建”时会收到错误消息。因此，请确保您在完成此步骤之前已收到批准消息。此外，如果您计划使用 us-west1 中的 NVIDIA T4 之外的其他 GPU 类型或位置，则需要在创建笔记本时更改这些参数。
您的实例需要几分钟才能启动，启动完成后，您将看到用蓝色字母显示的“打开 JUPYTERLAB”选项。
请注意，即使创建了 AI Platform Notebook 实例后，您也可以更改硬件（例如，添加或删除 GPU）。如果您将来需要这样做，只需停止实例，然后按照此处的步骤操作。
在 AI Platform Notebooks 上训练多 GPU 模型现在您的实例已设置好，您可以点击“打开 JUPYTERLAB”。
将 Colab 笔记本下载为 .ipynb 文件，并将其上传到您的 Jupyter Lab 环境。上传文件后，转到笔记本并运行代码。
执行model.fit单元格时，您应该注意到每个纪元的步骤数现在是 45，是使用单个 GPU 时的一半。这就是数据并行在起作用。在全局批次大小为 64 * 2 的情况下，您的 CPU 将向每个 GPU 发送 64 张图像的批次。因此，虽然以前模型在单个步骤中只看到 64 个示例，但现在它在每个步骤中看到 128 个示例，因此每个纪元所需的时间更少。以前每个纪元大约需要 100 秒，现在每个纪元大约需要 60 秒。您会注意到，添加第二个 GPU 并没有将时间缩短一半，因为同步梯度需要一些开销。对于更大的数据集（木薯只有 5656 个训练图像），好处会更加明显。此外，还有很多技巧可以帮助您从第二个 GPU 中获得更多收益，例如，确保您的输入管道不是瓶颈。要详细了解如何充分利用您的 GPU，请参阅TensorFlow 性能调试指南。
到目前为止，您已经了解了如何使用 GCP AI Platform Notebooks 运行简单的分布式训练作业。我们使用的数据集并不大，并且该模型在仅经过几个纪元后就达到了相当高的准确率。但是，在现实中，您的训练作业可能需要更长的时间运行，您可能不想使用笔记本。
启动 AI Platform Notebook 时，它会使用 GCP 深度学习虚拟机映像 创建一个 Google Compute Engine (GCE) 实例。深度学习虚拟机映像是针对数据科学和机器学习任务优化的 Compute Engine 虚拟机映像。在我们的示例中，我们使用了 TensorFlow Enterprise 2.3 映像，但是 还有许多其他选项可供选择。
在控制台中，您可以使用菜单导航到Compute Engine > 虚拟机实例
您应该看到一个与您之前创建的笔记本名称相同的实例。因为这是一个 GCE 实例，所以我们可以 ssh 到机器上并在那里运行代码。
安装 Google Cloud SDK 将使您能够从终端管理项目中的 GCE 资源。按照此处的步骤安装 SDK 并连接到您的项目。
SSH 到 VM安装和配置 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。
完成实验后，请务必关闭您的实例。您可以通过从控制台 UI 的笔记本实例页面或 GCE 实例页面中选择该实例，然后点击窗口顶部的“停止”来执行此操作。关闭实例非常重要，因为您将被收取每小时几美元的费用。您可以轻松地停止您的实例，然后在您想要运行更多实验时重新启动它，所有文件都将保留。
在本文中，您学习了如何使用 MirroredStrategy（一种同步数据并行策略）将 TensorFlow 训练作业分布到 GCP 上的两个 GPU 上。您现在了解了如何设置 GCP 环境和准备代码的基本机制，但分布式训练的世界还有很多东西值得探索。例如，如果您有兴趣将分布式训练作业构建到生产 ML 管道中，请查看AI Platform Training 服务，该服务还允许您跨多个机器（每个机器包含多个 GPU）配置训练作业。或者，您可以构建 TFX 管道，并在AI Platform Pipelines 服务上运行它。
在 tensorflow.org 网站上，您可以查看 概述指南中 tf.distribute.Strategy API 可用的其他策略，还可以了解如何在 自定义训练循环中使用策略。对于更高级的概念，有一个关于数据如何分布的指南，以及一个关于如何使用 TensorFlow Profiler 进行性能调试的指南，以确保您最大限度地利用 GPU。
