如何用Amazon SageMaker 做分布式 TensorFlow 训练？（千元亚马逊羊毛可薅）

机器之心转载，来源：AWS官方博客，作者：Ajay Vohra。

TensorFlow 是广泛被用于开发大型深度神经网络 (DNN) 的开放源机器学习 (ML) 库，此类 DNN 需要分布式训练，并且在多个主机上使用多个 GPU。Amazon SageMaker 是一项托管服务，可通过主动学习、超参数优化、模型分布式训练、监控训练进展，部署培训模型作为自动扩展的 RESTful 服务，以及对并发 ML 实验进行集中式管理，从标签数据开始简化 ML 工作流。

本文将重点讨论使用 Amazon SageMaker 进行分布式 TensorFlow 训练。

概念概述

尽管本文中的许多分布式训练概念通常适用于多种类型的 TensorFlow 模型，但本文着重要介绍的是在 Common Object in Context (COCO) 2017 数据集上对 Mask R-CNN 模型进行的分布式 TensorFlow 训练。

模型

Mask R-CNN 模型用于对象实例分割，其中模型生成像素级掩膜（Sigmoid 二进制分类）和以对象类别（SoftMax 分类）注释的目标框（Smooth L1 回归）描绘图像中的每个对象实例。Mask R-CNN 的一些常见用例包括：自动驾驶汽车的感知、表面缺陷检测和地理空间图像分析。

文中选择 Mask R-CNN 模型的关键原因有三个：

大型数据集上的 Mask R-CNN 分布式数据并行训练可增加通过训练管道的图像吞吐量，并缩短训练时间。
Mask R-CNN 模型有许多开源 TensorFlow 实现。本文使用 Tensorpack Mask/Faster-RCNN 实现作为主要示例，但也推荐使用高度优化的 AWS 示例 Mask-RCNN。
Mask R-CNN 模型在 MLPerf 结果中被评估为大型对象检测模型。

下图为 Mask R-CNN 深层神经网络架构的示意图。

分布式训练中的同步 Allreduce 梯度

分布式 DNN 训练的主要挑战在于，在应用梯度来更新跨多个节点的多个 GPU 上的模型权重之前，需要在同步步骤中对所有 GPU 的反向传播过程中计算出的梯度进行 Allreduce（平均化）。

同步 Allreduce 算法需要实现高效率，否则从分布式数据并行训练中获得的任何训练速度提升，都会因同步 Allreduce 步骤的效率低下而荡然无存。

要使同步 Allreduce 算法实现高效率，存在三个主要挑战：

该算法需要随着分布式训练集群中节点和 GPU 数量的增加而扩展。
该算法需要利用单个节点内的高速 GPU 到 GPU 互连的拓扑。
该算法需要通过有效地批处理与其他 GPU 的通信，来有效地交错 GPU 上的计算以及与其他 GPU 的通信。

Uber 的开放源库 Horovod 通过以下方式克服了这三个主要挑战：

Horovod 提供了一种高效的同步 Allreduce 算法，它可随着 GPU 和节点数量的增加而扩展。
Horovod 库利用 Nvidia Collective Communications Library (NCCL) 通信原语，而这些通信原语利用了对 Nvidia GPU 拓扑的了解。
Horovod 包含 Tensor Fusion，它通过批量处理 Allreduce 数据通信，来高效地将通信与计算交错。

许多 ML 框架（包括 TensorFlow）都支持 Horovod。TensorFlow 分发策略还利用了 NCCL，并提供了使用 Horovod 进行分布式 TensorFlow 训练的替代方法。本文使用的是 Horovod。

训练大型 DNN（如 Mask R-CNN）对每个 GPU 的内存要求较高，这样您才可以将一个或多个高分辨率图像推送经过训练管道。它们还需要高速 GPU 到 GPU 互连，以及高速网络互连机器，以便高效地对 Allreduce 梯度进行同步。Amazon SageMaker ml.p3.16xlarge 和 ml.p3dn.24xlarge 实例类型可以满足所有这些要求。如需更多信息，见 Amazon SageMaker ML 实例类型。它们具有八个 Nvidia Tesla V100 GPU、128–256 GB GPU 内存、25–100 Gbs 网络互连及高速 Nvidia NVLink GPU 到 GPU 互连，非常适合 Amazon SageMaker 上的分布式 TensorFlow 训练。

消息传递接口

分布式 TensorFlow 训练的下一个挑战是在多个节点上合理布置训练算法进程，以及将每个进程与唯一全局排名相关联。消息传递接口 (MPI) 是广泛用于并行计算的聚合通信协议，在管理跨多个节点的一组训练算法工作进程中非常有用。

MPI 用于在多个节点上布置训练算法进程，并将每个算法进程与唯一的全局和本地排名相关联。Horovod 用于逻辑地将给定节点上的算法进程固定到特定的 GPU。梯度同步 Allreduce 要求将每个算法进程逻辑固定到特定的 GPU。

在本文中，要理解的主要 MPI 概念是，MPI 在主节点上使用 mpirun，以便在多个节点上启动并发进程。主节点使用 MPI 管理着在多个节点集中运行的分布式训练进程的生命周期。要通过 MPI 使用 Amazon SageMaker 来进行分布式训练，您必须集成 MPI 和 Amazon SageMaker 的原生分布式训练功能。

集成 MPI 与 Amazon SageMaker 分布式训练

要了解如何集成 MPI 和 Amazon SageMaker 分布式训练，您需要对以下概念有相当认识：

Amazon SageMaker 要求训练算法和框架都打包在一个 Docker 映像中。
必须为 Amazon SageMaker 训练启用 Docker 映像。通过使用 Amazon SageMaker 容器可以简化启用，而该容器作为库则有助于创建已启用 Amazon SageMaker 的 Docker 映像。
您需要在 Amazon SageMaker 训练图像中提供入口点脚本（通常是 Python 脚本），以充当 Amazon SageMaker 和您的算法代码之间的中介。
要在指定主机上开始训练，Amazon SageMaker 会从训练图像运行一个 Docker 容器，然后使用提供信息（如超参数和输入数据位置）的入口点环境变量调用入口点脚本。
入口点脚本则使用在入口点环境变量中传递给它的信息启动具有正确 args 的算法程序，并对运行的算法进程进行轮询。
若算法进程退出，入口点脚本将使用算法进程的退出代码退出。Amazon SageMaker 使用此退出代码来确定训练作业成功与否。
入口点脚本会将算法进程的 stdout 和 stderr 重定向至它自己的 stdout。反过来，Amazon SageMaker 会捕获来自入口点脚本的 stdout，然后将其发送到 Amazon CloudWatch Logs。Amazon SageMaker 为训练作业中定义的算法指标解析 stdout 输出，然后将指标发送至 Amazon CloudWatch 指标。
当 Amazon SageMaker 启动要请求多个训练实例的训练作业时，它会创建一组主机，然后逻辑地将每个主机命名为algo-k，其中 k 是该主机的全局排名。例如，若训练作业请求四个训练实例，Amazon SageMaker 会把主机分别命名为 algo-1、algo-2、algo-3 和 algo-4。在网络上，主机可以使用这些主机名进行连接。

如果分布式训练使用 MPI，您需要一个在主节点（主机）上运行，而且控制着分布于多个节点（从 algo-1 到 algo-n，其中 n 为在您的 Amazon SageMaker 训练作业中请求的训练实例的数量）的全部算法进程生命周期的 mpirun 命令。不过，Amazon SageMaker 不会察觉 MPI，或任何您可能用于在多个节点分配算法进程的其他并行处理框架。Amazon SageMaker 将在运行于每个节点的 Docker 容器上调用入口点脚本。这就意味着，入口点脚本需要知道其节点的全局排名，并根据它在主节点或其他非主节点上是否被调用而执行不同的逻辑。

具体而言，对于 MPI，在主节点上被调用的入口点脚本需要运行 mpirun 命令，以开始当前 Amazon SageMaker 训练作业的主机集中全部节点的算法进程。在任何非主节点上被 Amazon SageMaker 调用时，同一个入口点脚本会定期检查由 mpirun 从主节点远程管理的非主节点上的算法进程是否依然在运行，并且在不运行时退出。

MPI 中的主节点是一个逻辑概念，它取决于入口点脚本在当前训练作业的全部主机中指定一个主机作为主节点。这项指定必须采用分散式方法完成。简单的做法之一是，将 algo-1 指定为主节点，而其他所有主机为非主节点。由于 Amazon SageMaker 在入口点环境变量中为每个节点提供其逻辑主机名，因此节点可以直观地确定其是主节点或非主节点。

包含于随附 GitHub 存储库并在 Tensorpack Mask/Faster-RCNN 算法 Docker 映像中打包的 train.py 遵循本部分中概述的逻辑。

如果具备这样的概念理解背景，您就可以继续操作分步教程，了解如何使用 Amazon SageMaker 为 Mask R-CNN 运行分布式 TensorFlow 训练。

解决方案概览

本教程有以下关键步骤：

使用 AWS CloudFormation 自动化脚本创建一个私有 Amazon VPC，以及一个附加于此私有 VPC 的 Amazon SageMaker 笔记本实例网络。
在附加于您的私有 VPC 的由 Amazon SageMaker 托管的 Amazon VPC 网络中，从 Amazon SageMaker 笔记本实例启动分布式训练作业。您可以使用 Amazon S3、Amazon EFS，和 Amazon FSx 作为训练数据管道的数据源。

先决条件

以下为必须满足的先决条件：

创建并激活一个 AWS 账户或使用现有的 AWS 账户。
管理您的 Amazon SageMaker 实例限制。您至少需要两个 ml.p3dn.24xlarge 或两个 ml.p3.16xlarge 实例，建议每个的服务限制为四个。记住，每个 AWS 区域都有特定的服务限制。本文使用的是 us-west-2。
克隆本文的 GitHub 存储库，并执行本文中的步骤。本文中的所有路径都相对于 GitHub 存储库根目录。
使用任何支持 Amazon SageMaker、EFS 和 Amazon FSx 的 AWS 区域。本文使用的是us-west-2。
创建一个新的 S3 存储桶或选择一个现有的。

创建附加于 VPC 的 Amazon SageMaker 笔记本实例

第一步是运行 AWS CloudFormation 自动化脚本以创建一个附加于私有 VPC 的 Amazon SageMaker 笔记本实例。要运行此脚本，您需要具有与网络管理员职能相符的 IAM 用户权限。如果没有此类权限，您可能需要寻求网络管理员的帮助以运行本教程中的 AWS CloudFormation 自动化脚本。如需更多信息，见工作职能的 AWS 托管策略。

使用 AWS CloudFormation 模板 cfn-sm.yaml 以创建一个 AWS CloudFormation 堆栈，而该堆栈将创建一个附加于私有 VPC 的笔记本实例。您可以使用 AWS CloudFormation 服务控制台中的 cfn-sm.yaml 以创建 AWS CloudFormation 堆栈，或者您也可以自定义 stack-sm.sh 脚本中的变量，并在您已安装 AWS CLI 的任何位置运行该脚本。

要使用 AWS CLI 方法，执行以下步骤：

安装 AWS CLI 并对其进行配置。
在 stack-sm.sh 中，将 AWS_REGION 和 S3_BUCKET 分别设为您的 AWS 区域和您的 S3 存储桶。您将要用到这两项变量。
或者，如果想要使用现有的 EFS 文件系统，您需要设置 EFS_ID 变量。如果您的 EFS_ID 留空，将创建一个新的 EFS 文件系统。若您选择使用现有的 EFS 文件系统，请确保现有的文件系统没有任何现有的挂载目标。如需更多信息，见管理 Amazon EFS 文件系统。
您还可以指定 GIT_URL 以添加 GitHub 存储库至 Amazon SageMaker 笔记本实例。如果是 GitHub 存储库，您可以指定 GIT_USER 和 GIT_TOKEN 变量。
运行自定义 stack-sm.sh 脚本以创建一个使用 AWS CLI 的 AWS CloudFormation 堆栈。

保存 AWS CloudFormation 脚本摘要输出以供稍后使用。您还可以在 AWS 管理控制台的 AWS CloudFormation 堆栈输出选项卡的下方查看输出。

启动 Amazon SageMaker 训练作业

在 Amazon SageMaker 控制台中，打开您创建的笔记本实例。在此笔记本实例中，有三个可用于训练 Mask R-CNN 的 Jupyter 笔记本：

Mask R-CNN 笔记本，它使用 S3 存储桶作为数据源：mask-rcnn-s3.ipynb。
Mask R-CNN 笔记本，它使用 EFS 文件系统作为数据源：mask-rcnn-efs.ipynb。
Mask R-CNN 笔记本，它使用 Amazon FSx Lustre 文件系统作为数据源：mask-rcnn-fsx.ipynb。

对于本文选用的 Mask R-CNN 模型和 COCO 2017 数据集，所有三个数据源选项的训练时间性能差不多（虽然不完全一样）。每个数据源的成本结构各有差异。以下是它们在设置训练数据管道的时间方面的差异：

对于 S3 数据源，在每次启动训练作业时，它将使用大约 20 分钟时间从您的 S3 存储桶复制 COCO 2017 数据集到附加于每个训练实例的存储卷。
对于 EFS 数据源，它将使用大约 46 分钟从您的 S3 存储桶复制 COCO 2017 数据集到您的 EFS 文件系统。您只需要复制此数据一次。在训练期间，将通过网络接口从挂载于所有训练实例的共享 EFS 文件系统输入数据。
对于 Amazon FSx，它将使用大约 10 分钟创建一个新的 Amazon FSx Lustre，并从您的 S3 存储桶将 COCO 2017 数据集导入到新的 Amazon FSx Lustre 文件系统。您仅需要执行此操作一次。在训练期间，将通过网络接口从挂载于所有训练实例的共享 Amazon FSx Lustre 文件系统输入数据。

如果不确定哪个数据源选项更适合您，您可以首先尝试使用 S3，如果每个训练作业一开始时的训练数据下载时间不可接受，接着探索并选择 EFS 或 Amazon FSx。不要对任何数据源的训练时间性能进行假设。训练时间性能取决于众多因素；最好的做法是进行试验与测量。

在所有三种情形中，训练期间的日志和模型检查点输出会被写入到附加于每个训练实例的存储卷，然后在训练完成时上传到您的 S3 存储桶。日志还会在训练过程中被注入到 Amazon CloudWatch，您可以在训练期间加以检查。系统和算法训练指标会在训练过程中被注入到 Amazon CloudWatch 指标，您可以在 Amazon SageMaker 服务控制台中对其进行可视化。

训练结果

下图为两种算法对 COCO 2017 数据集进行 24 次训练后的示例结果。

您可以在下方查看 TensorPack Mask/Faster-RCNN 算法的示例结果。下图可被拆分为三个存储桶：

不同并交比 (IoU)，以及小型、中型和大型对象大小值的目标框预测平均准确率 (mAP) 图示
不同并交比 (IoU)，以及小型、中型和大型对象大小值的对象实例分割 (segm) 预测平均准确率 (mAP) 图示
与训练损失或标签准确率有关的其他指标

您可以在下方查看优化 AWS Samples Mask R-CNN 算法的示例结果。下图显示的聚合 mAP 指标几乎与之前的算法相同，不过收敛进展各不相同。

结论

Amazon SageMaker 提供基于 Docker 的简化分布式 TensorFlow 训练平台，让您可以专注于您的 ML 算法而不会被从属的问题干扰，例如，基础设施可用性与可扩展性机制，以及并发实验管理，等等。在模型训练完成以后，您可以使用 Amazon SageMaker 的集成模型部署功能为您的模型创建一个自动可扩展的 RESTful 服务终端节点，并开始对其进行测试。如需更多信息，见在 Amazon SageMaker 托管服务上部署模型。若模型已准备就绪，您可以将模型 RESTful 服务无缝部署到生产。

关于作者

Ajay Vohra 是一名专门研究自动驾驶车辆开发感知机器学习的首席解决方案架构师。在加入 Amazon 前，Ajay 曾在一家本地数据中心从事金融风险建模的大规模并行网格计算，以及应用平台工程的自动化工作。