MXNet supports distributed training enabling us to leverage multiple machines for faster training.
MXNet支持分布式培训,使我们能够利用多台机器进行更快速的培训。这段话来自于 MXNet官网 ,说明了MXNet 支持跨越设备运行。
How to Start Distributed Training?
那么MXNet是怎样实现分布式训练的?
仔细阅读官方文档,官方文档给出了一个示例并写了这样一段话:
For distributed training of this example, we would do the following:
If the mxnet directory which contains the script image_classification.py is accessible to all machines in the cluster (for example if they are on a network file system), we can run:
../../tools/launch.py -n 3 -H hosts --launcher ssh python image_classification.py --dataset cifar10 --model vgg11 --epochs 1 --kvstore dist_sync
这是个使用 ssh 方式进行分布式训练的例子。
其中使用的测试代码,来源于
https://github.com/apache/incubator-mxnet
注意,这里的 launch.py 是一个非常重要的工具,如果仔细阅读 python 源码,会发现正是由于执行它才能实现分布式训练,它目前支持了5种分布式或并发训练方式:
–launcher denotes the mode of communication. The options are:
官网介绍了5种方式进行分布式或并行训练。这些方式都是以类似于集群管理的方式进行分布式训练。如果阅读整个测试代码,你可能还会发现其他几种集群管理方式,如mesos,不知道什么原因,官网没有介绍,也没有明确说支持 mesos 调度。
实际上,跟 tensorflow 或者 pytorch 进行分布式训练稍有不同。
在tensorflow 或者 pytorch 进行分布式训练,可能需要自己手动 或者 通过 mpi 工具起 不同角色,如 tensorflow 中的 ps 和 worker ,pytorch 中的 rank 。
而 MXNet 起不同的角色,全部都交给 launch.py 做完了。相当于对普通用户进行了一定程度上的屏蔽。
MXNet 这样做有一定的好处,普通用户只需要关注 训练脚本 的编写,而不需要关注 分布式计算 集群 如何运作。相对应的,当用户想抛开 launch.py 进行多机器多节点分布式训练 时,这也会成为弊病。因为MXNet 官网并没有仔细介绍如何手动启动 分布式训练。
export COMMAND='python example/gluon/image_classification.py --dataset cifar10 --model vgg11 --epochs 1 --kvstore dist_sync'
DMLC_ROLE=server DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=server DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=scheduler DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=worker DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=worker DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND
这是MXNet官方给出的在单节点机器,起一个 scheduler ,两个 worker, 两个 server 进行分布式训练。但是实际情况中,我们希望将 上述三种角色, 以在不同主机上进行启动,进行大规模硬件资源进行分布式训练。
另外,我们还希望通过仅仅提供硬件资源,以云计算的方式提供服务,用户之间具备隔离。也就是我们希望所有计算服务,以 docker 容器化方式启动。通过指定 worker 和 server 的数量以及要使用的资源,来启动容器。
如此一来,我们不能使用官方的 推荐 的 集群管理方式进行分布式训练。我们真正需要实现的是,launch.py 实现的功能,并以自动化的方式在分布式集群上 启动 docker 容器, 进行分布式训练。
所有的用户可以创建自己的 docker image,push 到分布式 image 仓库,并挂载不同的存储,达到用户隔离的目的。
这个 例子中用的 cifar10 数据 在执行过程中下载非常慢,我把它 放在https://download.csdn.net/download/github_37320188/11584523 , (由于没有积分了,小收几个积分用用,不能老做公益呀)
如果已经尝试过绕过 launch.py 进行分布式训练,最直接的方法是修改官方给出的 单机 并发训练 脚本。往上翻 可以看到。
实际上,该脚本通过设置环境变量的方式,设置 不同 角色的参数,然后执行 python 脚本。
官方对相关参数已经有部分解释。
但是尝试修改参数过后,如果成功找到方向的话,会发现,官方给出的环境变量是不够的。
这里提供 3 个解决思路去进行解决这个问题。
一开始我打印出以 ssh 方式分布式训练时 export 的环境变量,用来确定了确实是忽略了一些环境变量。
然后,我才开始阅读 launch.py 的代码,整个官方源码包其他的代码,找到了一些官方说明没有提到的变量。
操作一波
# 机器 1 输入
export COMMAND='python3 /mnt/mxnet-test/incubator-mxnet/example/gluon/image_classification.py --dataset cifar10 --model vgg11 --epochs 1 --kvstore dist_sync'
DMLC_ROLE=scheduler DMLC_PS_ROOT_URI=192.168.61.55 DMLC_PS_ROOT_PORT=9091 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=server DMLC_PS_ROOT_URI=192.168.61.55 DMLC_NODE_HOST=192.168.61.55 DMLC_PS_ROOT_PORT=9091 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 DMLC_SERVER_ID=0 $COMMAND &
DMLC_ROLE=worker DMLC_PS_ROOT_URI=192.168.61.55 DMLC_NODE_HOST=192.168.61.55 DMLC_PS_ROOT_PORT=9091 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 DMLC_WORKER_ID=0 $COMMAND
# 机器 2 输入
export COMMAND='python3 /mnt/mxnet-test/incubator-mxnet/example/gluon/image_classification.py --dataset cifar10 --model vgg11 --epochs 1 --kvstore dist_sync'
DMLC_ROLE=server DMLC_PS_ROOT_URI=192.168.61.55 DMLC_NODE_HOST=192.168.61.56 DMLC_PS_ROOT_PORT=9091 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 DMLC_SERVER_ID=1 $COMMAND &
DMLC_ROLE=worker DMLC_PS_ROOT_URI=192.168.61.55 DMLC_NODE_HOST=192.168.61.56 DMLC_PS_ROOT_PORT=9091 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 DMLC_WORKER_ID=1 $COMMAND
python 版本为 3.6.8
scheluder 不需要跟 任何其他 角色 绑定到 同一个 节点,也能独立运行。
甚至 可以把 这 5 个服务起在 5台机器上
DMLC_PS_ROOT_PORT 指定的是 起 scheduler 的机器上未被分配的可用端口
/mnt/mxnet-test/ 是一个分布式存储挂载目录
这里看到的是,相对于 官网 给到的介绍 多了几个环境变量:
如果你使用了ib 卡,除了直接指定 DMLC_NODE_HOST=ib卡的地址 ,还可通过 DMLC_INTERFACE 环境变量指定
除此之外,指定 DMLC_PS_ROOT_URI 和 DMLC_NODE_HOST 不能使用主机名,必须使用 ip 地址 或者 ib 卡的地址。原因是 MXNet 的通信依赖于 zmq ,其不支持 主机名。具体原因或者代码可以网上查阅
如果已经通过了上述 多机 分布式训练。后面就会变的稍微容易些。
首先,需要一个支持 python3 的 image,并且 支持 mxnet。你可以 在 docker hub 上找到一个
docker pull loongc/mxnet:v1
操作一波
# 机器 1 上的 启动 shell 脚本 mxnet-test.sh
docker run -d --env-file /tmp/mxnet_env/worker \
--name worker_1 \
-v /mnt/mxnet-test/incubator-mxnet:/incubator-mxnet \
-w /incubator-mxnet/example/gluon/ \
--net=host loongc/mxnet:v1 \
python3 image_classification.py --dataset cifar10 --model vgg11 --epochs 10 --kvstore dist_sync
docker run -d --env-file /tmp/mxnet_env/server \
--name server_1 \
-v /mnt/mxnet-test/incubator-mxnet:/incubator-mxnet \
-w /incubator-mxnet/example/gluon/ \
--net=host loongc/mxnet:v1 \
python3 image_classification.py --dataset cifar10 --model vgg11 --epochs 10 --kvstore dist_sync
docker run -d --env-file /tmp/mxnet_env/scheduler \
--name scheduler \
-v /mnt/mxnet-test/incubator-mxnet:/incubator-mxnet \
-w /incubator-mxnet/example/gluon/ \
--net=host loongc/mxnet:v1 \
python3 image_classification.py --dataset cifar10 --model vgg11 --epochs 10 --kvstore dist_sync
这里在机器 1 上启动了 3 个容器。
其中 的 /tmp/mxnet_env/worker , /tmp/mxnet_env/server , /tmp/mxnet_env/scheduler 文件 中写的是环境变量
这里举个例子
# 机器1 上的 scheduler 配置 /tmp/mxnet_env/scheduler
DMLC_ROLE=scheduler
DMLC_PS_ROOT_URI=192.168.61.55
DMLC_PS_ROOT_PORT=9091
DMLC_NUM_SERVER=2
DMLC_NUM_WORKER=2
# 机器1 上的 worker 配置 /tmp/mxnet_env/worker
DMLC_ROLE=worker
DMLC_PS_ROOT_URI=192.168.61.55
DMLC_PS_ROOT_PORT=9091
DMLC_WORKER_ID=0
DMLC_NODE_HOST=192.168.61.55
DMLC_NUM_SERVER=2
DMLC_NUM_WORKER=2
# 机器1 上的 server 配置 /tmp/mxnet_env/server
DMLC_ROLE=server
DMLC_PS_ROOT_URI=192.168.61.55
DMLC_PS_ROOT_PORT=9091
DMLC_NODE_HOST=192.168.61.55
DMLC_SERVER_ID=0
DMLC_NUM_SERVER=2
DMLC_NUM_WORKER=2
这些是机器 1 上的配置, 机器 2 上的配置很相似
区别在于:
- 机器 2 上的启动脚本 mxnet-test.sh 不需要启动 scheduler 容器,直接拿 机器 1 上脚本 改改 删掉 scheduler 部分
- 机器 2 上不需要 /tmp/mxnet_env/scheduler 配置脚本
- 机器 2 上 /tmp/mxnet_env/worker 和 /tmp/mxnet_env/server 两个文件中 DMLC_NODE_HOST 修改成机器2 的地址,DMLC_WORKER_ID、DMLC_SERVER_ID 修改为 1 (从 0 开始 计数)
分别启动两台机器上的 启动脚本 mxnet-test.sh ,无严格先后顺序。可以查看 /mnt/mxnet-test/incubator-mxnet/example/gluon/image-classification.log 查看执行情况。
其实,通篇内容主要是 多机多CPU 的分布式运算。当然,我们还做了 多机多GPU 的分布式运算,并且使用 ib 卡(InfiniBand)进行通信。
如果需要使用 多机多GPU 并且 基于Docker 进行分布式运算。你需要以下几个条件:
# 机器 1 上的 启动 shell 脚本 mxnet-test.sh
docker run -d --env-file /tmp/mxnet_env/worker \
--name worker_1 \
-v /mnt/mxnet-test/incubator-mxnet:/incubator-mxnet \
-w /incubator-mxnet/example/gluon/ \
--net=host mxnet/python:latest_gpu \
python3 image_classification.py --dataset cifar10 --model vgg11 --epochs 10 --kvstore dist_sync_device --gpus 0
docker run -d --env-file /tmp/mxnet_env/server \
--name server_1 \
-v /mnt/mxnet-test/incubator-mxnet:/incubator-mxnet \
-w /incubator-mxnet/example/gluon/ \
--net=host mxnet/python:latest_gpu \
python3 image_classification.py --dataset cifar10 --model vgg11 --epochs 10 --kvstore dist_sync_device --gpus 0
docker run -d --env-file /tmp/mxnet_env/scheduler \
--name scheduler \
-v /mnt/mxnet-test/incubator-mxnet:/incubator-mxnet \
-w /incubator-mxnet/example/gluon/ \
--net=host mxnet/python:latest_gpu \
python3 image_classification.py --dataset cifar10 --model vgg11 --epochs 10 --kvstore dist_sync_device --gpus 0
这里主要修改了 3 部分:
1)换了一个 image ,该 image 支持 gpu 运算
2)修改了 python3 命令。将 --kvstore 参数改为 dist_sync_device ,这是 gpu 支持的。 --kvstore 支持 4 个参数,详细内容在https://mxnet.incubator.apache.org/versions/master/faq/distributed_training.html 查看。
3)修改了 python3 命令。添加了 --gpus 0 参数。这是让 python 程序执行时使用 gpu ,如果使用两个 gpu,可写作–gpus 0,1 。具体内容也可以到上述页面查找。
另外,机器 2 上的 mxnet-test.sh 也需要做对应修改。
已经实现了手动分布式 训练,那么离 自动化 实现分布式训练就不困难了。
可以通过 k8s 或者 mesos 进行集群管理,通过调用 api 启动 指定数目的 worker 和 server 容器,并指定硬件资源,确定是否使用 ib 卡等。
如果有更好的实现 分布式训练 MXNet ,欢迎留言讨论。