MxNet系列——how_to——multi_devices

博客新址: http://blog.xuezhisd.top

邮箱：[email protected]

在多个CPU/GPUs上以数据并行方式运行MXNet

MXNet 支持在多个CPUs和GPUs上进行训练。其中，这些CPUs和GPUs可能位于不同的物理机上。

数据并行 vs 模型并行

MXNet模式使用数据并行的方式将工作负载划分到多个设备上。假如有 n 个设备，每一个设备都将获得完整的模型，并使用 1/n 的数据进行训练。结果（比如，梯度）和更新后的模型在不同设备之间进行通信。

MXNet也支持模型并行。这时，每个设备上维护模型的一部分。当模型非常大以至于单个设备无法存储时，模型并行非常有用。有一个关于如何在多层LSTM上使用模型并行方法的教程 a tutorial 。本教程主要关注 数据并行。

单机上的多个GPUs

划分工作负载

MXNet 默认会将一个数据批均匀地划分到每一个GPU上。加入批大小是 b，GPUs的数目是 k，因此在每一次迭代过程中，每一个GPU将会在 b/k 个样本上执行前向和后向传播。在更新模型之前，会将所有GPUs上的梯度都加起来。

如何使用

为了使用 GPUs，需要编译支持GPU的MXNet。

比如，在配置文件

config.mk

中设置

USE_CUDA=1

，然后再

make

。

(更多选项，请查看 MXNet installation guide)。

如果一台主机安装了一个或多个GPU，每个GPU都会有一个编号（编号从0开始计数）。如果想使用某个特定的显卡，既可以指定在代码中指定环境（context）

ctx

；也可以在命令行中传递参数

--gpus

。例如，如果想在Python中使用GPU 0和GPU 2，可以使用下面的代码创建网络模型。

import mxnet as mx
#列表包含多个GPU
model = mx.model.FeedForward(ctx=[mx.gpu(0), mx.gpu(2)], ...) 
           

如果程序接受参数

--gpus

，比如 example/image-classification，那么可以尝试下面的代码。

python train_mnist.py --gpus 0,2 ...
           

高级用法

如果多个GPUs的计算能力不同，那么可以根据它们的计算性能来划分工作负载。比如，如果GPU 0 是 GPU 2 性能的3倍，那么可以提供一个额外的负载选项

work_load_list=[3, 1]

。更多信息，请查看 model.fit。

如果所有其它超参都相同，在多个GPUs上训练结果应该和单GPU上的训练结果相同。但在实际应用中，由于随机存取（随机顺序或其它augmentations），使用不同的种子初始化权重和CuDNN，结果可能不同。

我们可以控制梯度聚合和模型更新（如果执行，训练过程）的位置，通过创建不同的

KVStore

（它是数据通信模块）。既可以使用

mx.kvstore.create(type)

来创建一个实例，也可以使用程序的参数

--kv-store type

来实现功能。

有两个常用的类型，

• local

  : 所有的梯度都复制到CPU内存，并且权重也在那里进行更新。

• device

  : 梯度聚集和权重更新都在GPU上进行。它也会尝试使用GPU的P2P通信（它可以加速通信）。但该选项会使用更多的GPU内存。

如果有大量的GPU（比如，>=4），我们建议使用

device

，以获得更佳性能。

使用多个机器进行分布式训练

我们可以通过简单地修改

KVStore

，实现在多个机器上运行MXNet。

• dist_sync

  。 它的行为和

  local

  相似。但一个主要区别是

  batch-size

  此处表示每个机器上的批大小。机器数量为 n，批大小为 b，

  dist_sync

  的行为等价于 批大小为 n/b

  local

  类型。

• dist_device_sync

  。它 和

  dist_sync

  之间的区别，与

  device

  和

  local

  之间的区别相同。即梯度聚集和权重更新都在GPU上进行。

• dist_async

  。它执行一步更新。 一旦从任何机器上接收到梯度，立即更新权重。这种更新是具有原子性（不可分），即同一权重在相同的时间不可能出现两个更新。然而却无法保证顺序。

如何启动一个作业

如果使用分布式训练，需要先设置

USE_DIST_KVSTORE=1

，再进行编译。

(更多选项，请查看 MXNet installation guide)。

启动一个分布式作业和单机运行稍有不同。MXNet 提供了 tools/launch.py ，它利用

ssh

,

mpi

,

sge

, 或

yarn

启动作业。

假定我们位于目录

mxnet/example/image-classification

下。希望在数据集mnist上使用脚本 train_mnist.py 训练LeNet。

单机运行时的命令：

python train_mnist.py --network lenet
           

现在有两个可以使用SSH进行通信的机器，我们希望在这两个机器上训练网络。首先，将这两个主机的IPs或主机名保存在

hosts

文件中。比如，

$ cat hosts
172.30.0.172
172.30.0.171
           

接下来，两个机器需要都可分访问mxnet文件夹，比如 网络文件系统。

运行以下命令，将在多机上启动MXNet。

../../tools/launch.py -n 2 --launcher ssh -H hosts python train_mnist.py --network lenet --kv-store dist_sync
           

注意：除了包括单机运行参数，此处还需：

• 使用

  launch.py

  来提交作业。

• --launcher

  提供启动器。如果所有的机器可以互相ssh，那么使用

  ssh

  ；如果

  mpirun

  可用，那么使用

  mpi

  ； 如果Sun公司的网格引擎（Sun Grid Engine）可用，使用

  sge

  ；如果Apache Yarn可用，使用

  yarn

  。

• -n

  计算节点（worker nodes）的数量

• -H

  指定host文件（

  ssh

  和

  mpi

  使用）

• --kv-store

  使用

  dist_sync

  或

  dist_async

同步目录

现在开始考虑 MXNet 文件夹不可访问的情况。我们可以先MXNet库复制到当前目录。

cp -r ../../python/mxnet .
cp -r ../../lib/libmxnet.so mxnet
           

然后，使用

launch.py

和参数

--sync-dst-dir

将当前目录同步到所有机器的

/tmp/mxnet

目录下。

../../tools/launch.py -n 2 -H hosts --sync-dst-dir /tmp/mxnet \
   python train_mnist.py --network lenet --kv-store dist_sync
           

使用一个特定的网络接口

MXNet 一般会选择第一个可用的网络接口。但对于有多个接口的机器，我们通过环境变量

DMLC_INTERFACE

可以指定使用哪个网络接口进行数据通信。例如，下面的代码中使用了网络接口

eth0

。

# 前面的命令设置网络接口（网卡）
# 后面的命令和上面相同
export DMLC_INTERFACE=eth0; ../../tools/launch.py ...
           

调试连接

通过设置

PS_VERBOSE=1

，可以查看调试日志。比如，

export PS_VERBOSE=1; ../../tools/launch.py ...
           

更多资料

• 通过命令

  ../../tools/launch.py -h

  查看更多启动选项。
• 查看 ps-lite 的更多选项。