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 支持分布式训练。它使我们能够利用多台机器进行更快、更有效的训练。

以下是两种方法,我们可以将训练 NN 的工作量分布到多个设备、CPU 或 GPU 设备 −

数据并行

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

模型并行

这是另一种并行,当模型太大而无法放入设备内存时,它就派上用场了。在模型并行中,不同的设备被分配学习模型不同部分的任务。这里需要注意的重点是,目前 Apache MXNet 仅支持单机模型并行。

分布式训练的工作原理

下面给出的概念是理解 Apache MXNet 中分布式训练工作原理的关键 −

进程类型

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

Worker

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

服务器

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

调度程序

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

KV 存储

KV 存储代表键值存储。它是用于多设备训练的关键组件。它很重要,因为单台机器上以及多台机器上跨设备的参数通信是通过一台或多台服务器传输的,其中的参数为 KVStore。让我们借助以下几点来理解 KVStore 的工作原理 −

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

  • 网络中的每个参数数组都分配有一个 key,该参数数组的权重由 value 引用。

  • 之后,工作节点在处理一批数据后 push 梯度。它们还会在处理新批次之前提取更新后的权重。

KVStore 服务器的概念仅在分布式训练期间存在,其分布式模式通过调用 mxnet.kvstore.create 函数来启用,该函数带有一个包含单词 dist

的字符串参数
kv = mxnet.kvstore.create('dist_sync')

密钥分发

所有服务器不必存储所有参数数组或密钥,但它们分布在不同的服务器上。这种密钥在不同服务器之间的分发由 KVStore 透明处理,并且随机决定哪个服务器存储特定密钥。

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

拆分训练数据

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

另一方面,在分布式训练的情况下,一开始我们需要将数据集分成n个不同的部分,以便每个工作器获得不同的部分。一旦得到,每个工作器就可以使用split_and_load再次将数据集的该部分划分到单个机器上的不同设备上。所有这些都是通过数据迭代器实现的。 mxnet.io.MNISTIteratormxnet.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 模式的工作方式相同,但属于异步模式。