博客新址: http://blog.xuezhisd.top
邮箱:[email protected]
MXNet 支持在多个CPUs和GPUs上进行训练。其中,这些CPUs和GPUs可能位于不同的物理机上。
MXNet模式使用数据并行的方式将工作负载划分到多个设备上。假如有 n 个设备,每一个设备都将获得完整的模型,并使用 1/n 的数据进行训练。结果(比如,梯度)和更新后的模型在不同设备之间进行通信。
MXNet也支持模型并行。这时,每个设备上维护模型的一部分。当模型非常大以至于单个设备无法存储时,模型并行非常有用。有一个关于如何在多层LSTM上使用模型并行方法的教程 a tutorial 。本教程主要关注 数据并行。
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
查看更多启动选项。