MXNet 结合 kubeflow 进行分布式训练

MXNet以数据并行的方式进行单机多卡训练

MXNet 支持在多CPU和GPU上进行训练。其中，这些CPU和GPU可以分布在不同的服务器上。

一台GPU机器上的每块GPU都有自己的编号(编号从0开始计数)。如果想使用某一块特定的GPU卡，即可以在代码中直接指定context(ctx); 也可以在命令行中传递参数--gpus。

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

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

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

我们可以控制梯度聚合和模型更新(如执行，训练过程)的位置，通过创建不同的KVStore（数据通信模块）。即可以使用mx.kvstore.create(type)来创建一个实例，也可以使用程序的参数--kv-store type来实现功能。

两种常用功能

MXNet以数据并行的方式进行多机多卡训练

在介绍MXNet分布式训练之前，先介绍几个关键性的概念方便理解MXNet的分布式训练:

三种进程类型

KVStore

KVStore是MXNet提供的一个分布式key-value存储，用来进行数据交换。KVStore本质的实现是基于参数服务器。

通过引擎来管理数据一致性,这使得参数服务器的实现变得简单，同时使得KVStore的运算可以无缝的与其他部分结合在一起。

使用两层的通讯结构，原理如下图所示。第一层的服务器管理单机内部的多个设备之前的通讯。第二层服务器则管理机器之间通过网络的通讯。第一层的服务器在与第二层通讯前可能合并设备之间的数据来降低网络带宽消费。同时考虑到机器内和外通讯带宽和延时的不同性，可以对其使用不同的一致性模型。例如第一层用强的一致性模型，而第二层则使用弱的一致性模型来减少同步开销。

通过调用mxnet.kvstore.create()函数并传递dist关键字参数来开启分布式训练的KVStore模式:

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

分布式训练模式

开启分布式训练

MXNet为了用户方便的进行分布式训练提供了一个tools/launch.py脚本。并且支持对多种类型的集群资源进行管理,如: ssh,mpirun,yarn,sge。

如上图所示，将代码clone到本地后，进入gluon目录，使用image_classification.py和CIFAR10数据集来对VGG11模型进行分布式训练。

虽然MXNet实现了多机多卡的分布式训练，但是在资源隔离，资源调度，资源限制以及大规模训练时数据共享都是不能满足需求的，所以接下来介绍下MXNet基于Kubeflow的大规模分布式训练。

MXNet结合kubeflow进行分布式训练

在将MXNet结合kubeflow进行分布式训练之前，首先需要安装kubeflow环境之前已经介绍了，这里就不在详细说明了。当kubeflow基础环境部署完成之后，需要针对MXNet安装mxnet-operator。

安装mxnet-operator

安装完成后，执行以下命令，检验mxnet是否安装成功：

kubectl get crd

输出如下内容代表mxnet-operator安装成功:

基于kubeflow测试MXNet分布式训练

示例代码:

https://github.com/deepinsight/insightface

Dockerfile文件内容:

由于我们是基于kubernetes的pv和pvc的方式使用数据卷，所有集群中需要事先安装好storage-class install，这样当用户创建pvc时，会通过storage-class自动的创建pv。

当创建好pv之后，用户可以将该数据卷mount到自己的开发机上，并将需要训练的数据集移到该数据卷。用于之后创建训练worker pod的时候，挂载到worker容器中，供训练模型使用。

kubectl create -f insightface-train.yaml

登录到具体的node计算节点通过docker logs命令查看训练的日志：

docker logs -f fc3d73161b27

虽然已经完成了mxnet结合kubeflow实现大规模的分布式训练，但是除了功能上的基本跑通，还存在很多因素影响分布式训练的性能，如: GPU服务器的网络带宽，普通的我们使用的以太网因为通信延迟的原因，会大大影响多机扩展性。InfiniBand（IB）网络和RoCE网络因为支持RDMA,大大降低了通信延迟，相比之下，20G的以太网格延迟会大大提升。当然，对于现有的普通以太网络，也可以通过别的方法优化通信带宽的减少，比方说梯度压缩。通过梯度压缩，减少通信带宽消耗的同时，保证收敛速度和精度不会有明显下降。MXNet官方提供了梯度压缩算法，按照官方数据，最佳的时候可以达到两倍的训练速度提升，同时收敛速度和精度的下降不会超过百分之一。还有如果使用分布式网络文件系统进行数据集的存储，如果解决吞吐量和网络延迟的问题。以及本地磁盘是否是SSD，还是在训练时是否需要对大文件的数据集进行record.io文件格式的处理及训练前数据集的切分等等问题，都需要更进一步的处理。

相关文章

• https://mxnet.apache.org/tutorials/vision/large_scale_classification.html

• https://www.kubeflow.org/docs/guides/components/mxnet/

• https://github.com/apache/incubator-mxnet

• https://github.com/apache/incubator-mxnet/issues/797

• https://transwarpio.github.io/teaching_ml/2016/07/05/mxnet/