Apache MXNet - 分布式训练

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

计算模式

MXNet 是一个多语言 ML 库，为用户提供以下两种计算模式 -

命令模式

这种计算模式公开了一个类似 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))

正如我们在上面的代码中看到的，MXNets 指定了保存张量的位置，无论是在 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 有以下三个进程 -

工人

Worker节点的工作是对一批训练样本进行训练。工作节点在处理每个批次之前将从服务器获取权重。处理完批次后，工作节点会将梯度发送到服务器。

服务器

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

调度程序

调度器的作用是建立集群，其中包括等待每个节点已经上来的消息以及节点正在监听哪个端口。设置集群后，调度程序让所有进程了解集群中的每个其他节点。这是因为进程可以相互通信。只有一个调度程序。

KV商店

KV 存储代表键值存储。它是用于多设备培训的关键组件。这很重要，因为单台以及多台机器上的设备之间的参数通信是通过一台或多台带有参数 KVStore 的服务器进行传输的。让我们借助以下几点来了解 KVStore 的工作原理 -

  • KVStore 中的每个值都由一个键和一个值表示。

  • 网络中的每个参数数组都分配有一个键，并且该参数数组的权重由值引用。

  • 之后，工作节点在处理完一批后推送梯度。他们还在处理新批次之前提取更新的权重。

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

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

密钥分配

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

如上所述，KVStore 确保每当提取密钥时，其请求都会发送到具有相应值的服务器。如果某个key的值很大怎么办？在这种情况下，它可以在不同的服务器之间共享。

分割训练数据

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

另一方面，在分布式训练的情况下，一开始我们需要将数据集分为n 个不同的部分，以便每个工作人员获得不同的部分。一旦获得，每个工作人员就可以使用split_and_load再次将数据集的该部分划分到一台机器上的不同设备上。所有这一切都是通过数据迭代器发生的。mxnet.io.MNISTIterator和mxnet.io.ImageRecordIter是 MXNet 中支持此功能的两个迭代器。

权重更新

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

  • 第一种方法聚合梯度并使用这些梯度更新权重。

  • 在第二种方法中，服务器仅聚合梯度。

如果您使用的是 Gluon，可以通过传递update_on_kvstore变量来选择上述方法。让我们通过创建训练器对象来理解它，如下所示 -

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 启用的不同分布式训练模式 -

距离同步

顾名思义，它表示同步分布式训练。在这种情况下，所有工作人员在每个批次开始时都使用相同的同步模型参数集。

这种模式的缺点是，在每个批次之后，服务器必须等待接收来自每个工作人员的梯度，然后才能更新模型参数。这意味着如果一个worker崩溃，所有worker的进度都会停止。

分布式异步

顾名思义，它表示同步分布式训练。在这种情况下，服务器从一名工作人员接收梯度并立即更新其存储。服务器使用更新的存储来响应任何进一步的拉取。

与dist_sync 模式相比，优点是处理完一个批次的工作人员可以从服务器拉取当前参数并开始下一个批次。即使其他工作人员尚未完成处理较早的批次，该工作人员也可以这样做。它也比 dist_sync 模式更快，因为它可能需要更多的纪元才能收敛，而无需任何同步成本。

距离同步设备

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

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

远程异步设备

此模式的工作方式与dist_sync_device模式相同，但采用异步模式。