数据集暴增压力下，微信「扫一扫」识物训练如何优雅破局？

微信“扫一扫”识物上线一段时间，由前期主要以商品图（鞋子/箱包/美妆/服装/家电/玩具/图书/食品/珠宝/家具/其余）做为媒介来挖掘微信内容生态中有价值的信息，扩张到各类垂类领域的识别，包括植物/动物/汽车/果蔬/酒标/菜品/地标识别等，识别核心依托于深度学习的卷积神经网络模型。随着天天千万级的增加数据和愈来愈多的模型参数量，深度学习训练一次时间大概须要一周左右。如何可以快速训练优化模型并上线，成为咱们亟待解决的问题。

1、引言

现在，依托强大的GPU算力，深度学习获得迅猛发展。在图像处理、语音识别领域掀起了史无前例的一场革命。相较于传统的方法，以卷积神经网络（CNN）为表明的深度学习方法能够高度地重点学习数据的特性，在图像处理领域已经取得了统治地位。

随着扫一扫识物日调用量的持续增长，图片数据正以天天千万级的量级增加，在这个争分夺秒的时代里，得数据者得天下。同时，神经网络的复杂性呈爆炸式增加，像15年微软提出图像分类的ResNet模型有7 ExaFLOPs/6千万个参数，17年谷歌的神经网络机器翻译模型有100 ExaFLOPS/87亿个参数。

在大部分场景下，模型能够在一台GPU服务器上，使用一个或者多个GPU进行训练。但随着数据集的增大，训练时间也相应增加，有些时候训练须要一周甚至更长时间。所以，如何可以快速迭代优化深度学习模型，成为咱们算法开发者亟须解决的问题。

下文将经过从分布式训练方法的选择、多机通讯技术原理进行讲解，基于Horovod的框架在微信自研平台打通分布式训练和实验结果来介绍微信扫一扫识物中的深度学习模型分布式训练。

2、分布式训练

1. 并行方式

多机多卡相比较于单机单卡，能够将模型训练的时间大大缩短。通常咱们一台服务器只支持8张GPU卡，而采用分布式的多机多卡训练方式，能够将几十甚至几百台服务器调度起来一块儿训练一个模型，进一步突破模型训练的上限。

按照分布式并行训练方式，分布式训练通常分为数据并行和模型并行两种。

（1）数据并行

分布式系统中不一样的GPU都有同一个模型的彻底拷贝，每一个GPU只得到整个数据的不一样部分，而后将全部GPU的神经网络梯度按照同步或者异步的方式合并。

（2）模型并行

分布式系统中每一个GPU使用相同的数据，全部GPU上只分布着模型的部份内容，训练过程当中交换神经网络的激活部分。

由于模型并行各个部分存在必定的依赖，不能随意增长GPU的数量，规模伸缩性差，在实际训练中用的很少。而数据并行，各部分独立，规模伸缩性好，实际训练中更为经常使用，提速效果也更好。在实现性、容错性和好的集群利用率上，数据并行优于模型并行。

2. 系统架构

分布式训练系统架构主要包括两种：Parameter Server Architecture（就是常见的PS架构，参数服务器）和Ring all-reduce Architecture。

（1）Parameter Server 架构

在PS架构中，集群中的节点被分为两类：parameter server和worker。其中parameter server存放模型的参数，而worker负责计算参数的梯度。

在每一个迭代过程，worker从parameter sever中得到参数，而后将计算的梯度返回给parameter server，parameter server聚合从worker传回的梯度，而后更新参数，并将新的参数广播给worker。

（2）Ring all-reduce 架构

在Ring all-reduce架构中，各个设备都是worker，而且造成一个环，没有中心节点来聚合全部worker计算的梯度。在一个迭代过程，每一个worker完成本身的mini-batch训练，计算出梯度，并将梯度传递给环中的下一个worker，同时它也接收从上一个worker的梯度。对于一个包含N个worker的环，各个worker须要收到其它N-1个worker的梯度后就能够更新模型参数。

采用PS计算模型的分布式，一般会遇到网络的问题，随着worker数量的增长，其加速比会迅速的恶化。相比PS架构，Ring all-reduce架构网络通讯量不随着worker（GPU）的增长而增长，是一个恒定值，集群中每一个节点的带宽都被充分利用。

3. 参数更新

（1）同步更新

全部 GPU 在同一时间点与参数服务器交换、融合梯度。在每轮训练的时候须要汇总全部 worker训练获得的梯度值，而后取平均值来更新参数服务器上的模型参数。

（2）异步更新

全部GPU 各自独立与参数服务器通讯，交换、融合梯度。每一个 worker 在每轮训练开始前从参数服务器获取模型参数，读取训练数据，进行训练，训练结束后便当即应用梯度来更新参数服务器上的模型参数。

异步更新通讯效率高速度快，但每每收敛不佳，由于一些速度慢的节点总会提供过期、错误的梯度方向。同步更新通讯效率低，一般训练慢，但训练收敛稳定，由于同步更新基本等同于单卡调大的batch size训练。可是传统的同步更新方法（各个GPU卡算好梯度，求和算平均的方式），在融合梯度时，会产生巨大的通讯数据量。

经过比对不一样分布式并行方式、系统架构和参数更新，微信扫一扫识物最终选择基于数据并行的参数同步更新的Ring all-reduce的分布式训练方法。

3、多机通讯技术

相比于单机多卡，多机多卡分布式训练要保证多台机器之间是能够互相通讯的以及不一样机器之间梯度可传递。

并行任务的通讯通常能够分为点对点通讯和集体通讯。点对点通讯这种模式只有一个sender和一个receiver，实现起来比较简单。而涉及到分布式训练，通常是多台服务器，用到集体通讯模式，包含多个sender多个receiver。集体通讯经常使用的通讯方式主要有下面几个：broadcast、gather、scatter、reduce、all-reduce等。

1. MPI

在微信的自研训练平台中，多机的通讯是基于消息传递接口(Message Passing Interface，MPI)来实现的，MPI是一种基于信息传递的并行编程技术，定义了一组具备可移植性的编程接口，是一种编程接口标准。

在基于MPI编程模型中，计算是由一个或多个彼此经过调用库函数进行消息收、发通讯的进程所组成。MPI中的通信器定义了一组可以互相发消息的进程。在这组进程中，每一个进程会被分配一个序号，称做秩(rank)，进程间显性地经过指定秩来进行通讯。MPI涉及到的一些操做包括数据移动，汇集、同步等。

因为深度学习训练参数大多在GPU上的，若是只是依靠MPI来同步参数，参数须要先从GPU搬到CPU，而后不一样机器CPU之间再通讯，通讯结束以后再将参数从CPU搬到GPU，这个过程的通讯效率是很低的。因此为了提升通讯效率，在训练的过程当中使用基于nvidia开发的NCCL进行通讯。

2. NCCL

NCCL是Nvidia Collective multi-GPU Communication Library的简称，是Nvidia开发的可以实现多GPU的集体通讯的库，可以很方便的集成至任何深度学习的训练框架。在实现 Allreduce、Reduce、Broadcast、Allgather等方面作了不少优化，能够在PCIe、Nvlink、InfiniBand上实现较高的通讯速度。

目前NCCL1.0版本只支持单机多卡，卡之间经过PCIe、NVlink、GPU Direct P2P来通讯。NCCL 2.0会支持多机多卡，多机间经过Sockets (Ethernet)或者InfiniBand with GPU Direct RDMA通讯。

4、Horovod训练框架

目前分布式训练框架有许多，Horovod 是 Uber 开源的一个深度学习工具，囊括了TensorFlow, Keras, PyTorch, and Apache MXNet 这些分布式训练框架。

而且Horovod的梯度同步和权值同步利用基于MPI和NCCL的 all-reduce算法,而非参数服务器架构,通讯效率更高。Horovod能够利用NVLINK、RDMA、GPUDirectRDMA、自动检测通讯拓扑以及回退到 PCIe 和 TCP/IP 通讯这些功能。同时，将已有的训练代码改为分布式训练代码，改动量少，简化分布式训练的运行和启动。

基于此，微信扫一扫识物选择Horovod的分布式训练框架，在微信自研的训练平台上进行训练。

Horovod的多机通讯初始化是基于MPI的,经过MPI初始化通讯环境和进程分配。有以下几个经常使用的环境参数：

• size: 进程数量，也就是GPU数量；
• rank：进程的惟一ID， 0-size；
• local size: 每一个worker的本地进程数量；
• local rank: 每一个worker的进程本地惟一ID。

经过这些参数来控制机器进程之间的通讯。

因为训练采用的是数据并行这种模式，因此须要对数据进行分布式采样。Horovod能够直接调用pytorch自带的分布式采样函数torch.utils.data.distributed.DistributedSampler。

这种方式能够适用于简单的分布式训练任务。可是在识物的检索训练过程当中，咱们但愿dataloader能够作一些平衡采样或者三元组采样，上面的sampler只支持分布式采样。

因为pytorch的DataLoader的部分初始化参数之间存在互斥关系，若是自定义了sampler,那么这些参数batch_size、shuffle、batch_sampler、drop_last都必须使用默认值。因此咱们重写batch_sampler，将分布式的sampler做为参数传入新构造的batch_sampler。

Horovod内部实现了广播操做，使模型在全部工做进程中实现一致性初始化。在加载模型权重的时候，只要在rank0的机器上加载，而后使用广播机制，就能够将参数同步到其余机器上进行权重初始化。

在训练过程当中，计算损失函数时须要涉及到allreduce操做，将全部worker的损失规约，而后再进行梯度传播。最后在保存模型时，只要指定一台机器保存模型便可。

5、实验结果

分布式训练除了训练阶段的通讯要尽量的快，数据的IO也是须要考虑的地方。扫一扫识物的检索模型是基于大量的图像数据训练的。在进行分布式训练时，每一个机器都须要可以读取这些训练数据，图片文件存到微信自研分布式存储系统上。

在训练时，分布式训练的加速比和GPU数目正相关。在mnist数据集上基于resnet50测试分布式训练运行时间, 单机运行mnist 100个epoch须要78min，多机使用4块GPU训练100个epoch 23min。

在咱们实际项目的模型训练中，基于分布式训练能够将以往须要训练5天甚至一周的时间缩短到1天之内，在一样的时间内，算法开发者能够探索更多的实验，快速反馈更新，大大提升了算法研发的效率。

6、总结与展望

目前扫一扫识物在微信自研训练平台上可以成功进行分布式训练，但仍然存在如下问题：如何可以高效地存放读取大量图片小文件，减小IO的耗时。道阻且长,行则将至，在后续工做中咱们将针对这些问题进行探索。

参考文献：

[1] Li M, Andersen D G, Park J W, et al. Scaling distributed machine learning with the parameter server[C]//11th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 14). 2014: 583-598.

[2]https://mpitutorial.com/tutor...

[3]https://developer.nvidia.com/...

[4]https://eng.uber.com/horovod/