Apache MXNet - 分布式训练

本章介绍 Apache MXNet 中的分布式训练。让我们首先了解 MXNet 中的计算模式。

计算模式

MXNet 是一款多语言机器学习库，为用户提供了以下两种计算模式：

命令式模式

这种计算模式提供类似 NumPy API 的接口。例如，在 MXNet 中，使用以下命令式代码在 CPU 和 GPU 上构建一个全零张量：

import mxnet as mx
tensor_cpu = mx.nd.zeros((100,), ctx=mx.cpu())
tensor_gpu= mx.nd.zeros((100,), ctx=mx.gpu(0))

如上代码所示，MXNet 指定了张量存储的位置，可以在 CPU 或 GPU 设备上。在上面的例子中，位置为 0。MXNet 达到了惊人的设备利用率，因为所有计算都是延迟执行的，而不是立即执行。

符号式模式

虽然命令式模式非常有用，但这种模式的一个缺点是其僵化性，即所有计算都需要预先知道，并且需要预定义的数据结构。

另一方面，符号式模式提供类似 TensorFlow 的计算图。它通过允许 MXNet 使用符号或变量而不是固定/预定义的数据结构来克服命令式 API 的缺点。之后，这些符号可以解释为一组操作，如下所示：

import mxnet as mx
x = mx.sym.Variable(“X”)
y = mx.sym.Variable(“Y”)
z = (x+y)
m = z/100

并行类型的

Apache MXNet 支持分布式训练。它使我们能够利用多台机器进行更快、更有效的训练。

以下是将神经网络训练的工作负载分布到多个设备（CPU 或 GPU 设备）的两种方法：

数据并行

在这种并行性中，每个设备存储模型的完整副本，并使用数据集的不同部分。设备还会集体更新共享模型。我们可以将所有设备放在一台机器上，也可以跨多台机器。

模型并行

这是另一种并行性，当模型太大而无法放入设备内存时非常有用。在模型并行中，不同的设备被分配学习模型的不同部分的任务。这里需要注意的重要一点是，目前 Apache MXNet 只支持单机模型并行。

分布式训练的工作原理

以下概念是理解 Apache MXNet 中分布式训练工作原理的关键：

进程类型

进程之间相互通信以完成模型的训练。Apache MXNet 有以下三个进程：

工作节点

工作节点的任务是在一批训练样本上执行训练。工作节点在处理每一批之前会从服务器拉取权重。工作节点在处理完一批后会将梯度发送到服务器。

服务器

MXNet 可以有多个服务器来存储模型的参数并与工作节点通信。

调度器

调度器的作用是设置集群，包括等待每个节点启动的消息以及节点正在监听的端口。设置集群后，调度器让所有进程都知道集群中的其他每个节点。这是因为进程可以相互通信。只有一个调度器。

KVStore

KVStore 代表**键值**存储。它是用于多设备训练的关键组件。它很重要，因为在单机或跨多台机器的设备之间进行参数通信是通过一个或多个具有参数 KVStore 的服务器传输的。让我们通过以下几点来了解 KVStore 的工作原理：

KVStore 中的每个值都由一个**键**和一个**值**表示。
网络中的每个参数数组都分配一个**键**，并且该参数数组的权重由**值**引用。
之后，工作节点在处理完一批后**推送**梯度。它们还在处理新批次之前**拉取**更新的权重。

KVStore 服务器的概念仅在分布式训练期间存在，并且其分布式模式是通过使用包含单词**dist** 的字符串参数调用**mxnet.kvstore.create** 函数来启用的：

kv = mxnet.kvstore.create(‘dist_sync’)

密钥的分布

并非所有服务器都存储所有参数数组或键，而是将它们分布在不同的服务器上。KVStore 透明地处理这种跨不同服务器的键的分布，并且随机决定哪个服务器存储特定键。

如上所述，KVStore 确保每当拉取键时，其请求都会发送到拥有相应值的服务器。如果某个键的值很大怎么办？在这种情况下，它可能被共享到不同的服务器上。

分割训练数据

作为用户，我们希望每台机器都能处理数据集的不同部分，尤其是在数据并行模式下运行分布式训练时。我们知道，为了分割数据迭代器提供的一批样本以便在单个工作节点上进行数据并行训练，我们可以使用**mxnet.gluon.utils.split_and_load**，然后将批次的每一部分加载到将进一步处理它的设备上。

另一方面，对于分布式训练，首先我们需要将数据集分成**n**个不同的部分，以便每个工作节点获得不同的部分。一旦获得，每个工作节点可以使用**split_and_load**再次将数据集的那一部分划分到单机上的不同设备上。所有这些都是通过数据迭代器完成的。**mxnet.io.MNISTIterator** 和**mxnet.io.ImageRecordIter** 是 MXNet 中支持此功能的两个这样的迭代器。

权重更新

为了更新权重，KVStore 支持以下两种模式：

第一种方法聚合梯度并使用这些梯度更新权重。
在第二种方法中，服务器只聚合梯度。

如果使用 Gluon，可以通过传递**update_on_kvstore**变量来在这两种方法之间进行选择。让我们通过创建**trainer**对象来了解它：

trainer = gluon.Trainer(net.collect_params(), optimizer='sgd',
   optimizer_params={'learning_rate': opt.lr,
      'wd': opt.wd,
      'momentum': opt.momentum,
      'multi_precision': True},
      kvstore=kv,
   update_on_kvstore=True)

分布式训练模式

如果 KVStore 创建字符串包含单词 dist，则表示已启用分布式训练。以下是通过使用不同类型的 KVStore 可以启用的不同分布式训练模式：

dist_sync

顾名思义，它表示同步分布式训练。在此模式下，所有工作节点在每批的开始都使用相同的同步模型参数集。

此模式的缺点是，在每批之后，服务器必须等待接收来自每个工作节点的梯度，然后才能更新模型参数。这意味着如果一个工作节点崩溃，它将阻止所有工作节点的进度。

dist_async

顾名思义，它表示异步分布式训练。在此模式下，服务器接收来自一个工作节点的梯度并立即更新其存储。服务器使用更新后的存储来响应任何进一步的拉取请求。

与**dist_sync 模式**相比，它的优势在于，完成处理一批的工作节点可以从服务器拉取当前参数并开始下一批。即使其他工作节点尚未完成处理之前的批次，工作节点也可以这样做。它也比 dist_sync 模式快，因为它可以在没有任何同步成本的情况下花费更多轮次来收敛。

dist_sync_device

此模式与**dist_sync**模式相同。唯一的区别是，当每个节点上使用多个 GPU 时，**dist_sync_device** 在 GPU 上聚合梯度并更新权重，而**dist_sync** 在 CPU 内存上聚合梯度并更新权重。

它减少了 GPU 和 CPU 之间的昂贵通信。这就是为什么它比**dist_sync**更快的原因。缺点是它会增加 GPU 的内存使用量。

dist_async_device

此模式的工作方式与**dist_sync_device**模式相同，但在异步模式下。

打印页面