深度学习的分布式训练

在深度学习这一领域经常涉及到模型的分布式训练（包括一机多GPU的情况）。我自己在刚刚接触到一机多卡，或者分布式训练时曾对其中到底发生了什么有过很多疑问，后来查看了很多资料，在这篇博客里对分布式的深度学习模型训练做一个总结。

由于是我自己的总结，所以如果有说的不对的地方，还望指正。

当然我的总结也是基于一些文献参考来的，所以如有雷同，纯属我借鉴别的大神的文章。

如果去查阅与分布式深度学习相关的资料，一般会看到两个词，模型并行和数据并行。

模型并行是指把模型的不同部分放置在各个设备上，这样整个集群就像一个生产流水线一样，对这样的分布式设置我还没有使用过，所以不在这里细说。

数据并行应该是大多数人都会遇到的场景。数据并行分为同步更新和异步更新，下面一一说明。

深度学习训练的数据并行化，说白了就是梯度下降的并行化，所以想要理解你在并行化训练一个深度学习的模型时到底发生了说明，你必须明白梯度下降时到底发生了什么。

我们先从单机的梯度下降说起。

所有介绍梯度下降的参考书都会说梯度下降分为整体梯度下降（每次更新参数使用所有数据）和随机梯度下降（每次更新参数使用一条数据），而在深度学习里我们一般都会结合这两种思想使用批梯度下降，每次更新参数使用一个batch的数据来更新参数。

一般我们可以把随机梯度下降的公式写成下面的形式：

这里我没有在θ上加任何标记，代指神经网络中随意的一个参数，这个公式搞神经网络的应该很熟悉，这里α是学习率，J(θ)是损失函数，整个式子的意义就是参数θ的更新是用上一次迭代的θ减去学习率乘以损失函数对该θ的偏导，也就是梯度。计算该偏导数的时候，就会用到我们的样本数据，不然无法计算J(θ)对θ的偏导数。

而深度学习常用的批梯度下降，是指每次更新时使用一个batch的数据量去更新参数，这句话的意思是说上面的公式中：

这个偏导数的部分，也就是梯度部分，是用这一个batch每一条数据计算出的梯度取平均数后再带入公式计算。所以说这是一种结合了随机梯度下降和整体梯度下降的方法。这么做的想法来源其实就是用抽样出来的样本的平均数来代替整体的平均数。

理解上面这段话很重要，如果不理解这句话就没办法理解分布式梯度下降到底发生了什么。

下面来说明分布式训练时发生的事，这得分为同步更新和异步更新，假设我们拥有这样一个集群：

参数服务器（PS）1台，记为ps，worker10台，每台上面都有一块GPU，记为w1，w2……w10，而我们的数据自然也被分成了10份，每台worker上一份，每一份都是不重复的，假设我们把batchsize设为128，并假设我们的模型一共有n个参数，再假设我们是在tensorflow上做模型的训练。

那么参数更新的过程大体如下：

参数服务器收集每台worker计算的梯度，更新参数，然后把更新后的参数再分发到各台worker，接着各台worker再根据新的参数计算梯度。

同步更新：

同步更新的每一次迭代，每台worker计算一个batchsize的梯度，在这里对每一个参数就会生成一个128维的列向量，每一维代表根据一条数据计算出的梯度，所以最终会生成一个128xn的矩阵（或者是nx128，这得看怎么实现的）。接着worker把这个矩阵传递给参数服务器，参数服务器等待所有的10台worker把梯度矩阵都传过来后（此时参数服务器上对每一个参数，都收到了1280个梯度值），便会对每个参数的1280个梯度值做一个平均，做完平均就可以代入上面的梯度更新的公式更新参数。更新完参数再把新的参数值传输回每台worker，每台worker再根据新的参数值继续计算梯度。

根据以上描述，我们可以得到关于同步更新的下面几个事实：

1. 虽然我们在tensorflow的代码里设置了batchsize为128，但其实模型在真正更新的时候是用了一个1280的大batch在更新参数
2. 同步更新的信息传输开销很大，而且有短板效应，如果其中一台worker的GPU比其他几台的差，那么每次更新时都需要等待这台worker完成计算
3. 同步更新并没有加快每次迭代的速度，相反，由于巨大的通信开销，每次迭代的时间可能比单机模型下还要慢
4. 同步更新时，整个模型更新的global_step和每台worker的local_step是相等的，如果我们设置最大步数为20000步，那么每台worker都会跑20000步。

同步更新是一种很容易想到的数据并行分布式梯度下降的过程，优点是模型的收敛会比较平稳，因为这种方式用了一个很大的batch，batch越大那么这个批梯度下降的效果就会越接近整体梯度下降。但同步更新的缺点就是信息传输开销很大，并且有短板效应。

为了避免短板效应，tensorflow中做了这样的设置，我们现在有10台worker，你可以指定ps只要接收到8份梯度，就开始做平均和参数更新，那么这样就可以避免每次更新都等待最慢的那台worker。

对于每次迭代的时间可能比单机还要慢这种说法，这是我在使用tensorflow时的切身感受，当然这慢多少和网络带宽是有关系的。但每次迭代慢并不意味着最终的收敛时间慢，因为你使用了更大的batch，所以模型可能会更快的收敛。也就是说本来你达到99%的准确率需要训练20000步，现在可能只需要10000步，并且有可能会继续收敛。

对于我个人的使用而言，我觉得同步更新解决了我这样的一个痛点：我可以设置更大的batchsize，因为batchsize越大，每次参数更新时保存的梯度矩阵就会越大，有可能会超出GPU的显存，这样如果你用的GPU显存不大，那么一块GPU能承受的batchsize就是有限的，而多GPU并行后，使用同步更新时batchsize就会自然而然的扩大（实际代码里的batchsize没变，但计算时其实变大了）。

异步更新：

异步更新的每一次迭代，每台worker计算好每个参数的128个梯度后，便会把梯度矩阵传输到ps上，ps在接受到任意一台worker的梯度矩阵后，就会立即更新参数，并把更新后的参数分发回worker，接着worker再根据新的参数抽取128条数据计算梯度。

根据上面的描述，我们可以得到关于异步更新的下面几个事实：

1. 异步更新时ps不需要再等待全部的10台worker都计算好梯度后再更新
2. 异步更新时，我们在代码里写的batchsize为128，那么每次更新时确确实实是用的128条数据的平均梯度在更新参数，而不是像同步更新时是1280条数据
3. 异步更新时，整体模型更新的global_step等于各台worker的local_step之和。比如我们设置总的迭代次数，也就是global_stpe为20000步，那么每台worker只要跑2000步迭代就会停止，因为每台worker的每一步local_step都会让参数更新一次，10台worker的话每台worker只要跑2000步。这也是和同步更新不一样的地方。

乍一看，异步更新好像是就是单机版的更新方式，batchsize和单机时一样，全部的步数也和设置的一样，但异步更新时会遇到过期梯度的问题。

过期梯度可以这么理解，比如说现在模型的global_step为100，记这时的参数为100版本的参数，worker1拿到第100版本的参数后开始计算梯度，而这时worker2刚刚把自己计算的梯度传到ps，那么ps就会立即更新参数到101版本，过了一会worker1还在计算梯度，这时worker3也计算好梯度了，于是ps立马又把参数更新到102版本，这样几个过程后，假设等worker1根据100版本的参数计算好梯度后，ps上的参数版本已经是102版本了，那么ps就会在102参数的版本基础上，根据worker1传来的梯度进行更新，也就是利用worker1根据100版本的参数计算出的梯度，把参数从102版本更新到103版本，就像下面的公式这样：

上面更新的公式里，θ的上标表示参数的版本（global_step），上文说的worker1的情况就和展示出来的这三道公式类似，前面两道更新对应worker2，worker3更新，最后一道对应worker1的更新。

对于出现这种过期梯度的情况，会导致梯度下降的过程变得不稳定，因为模型在更新参数时使用的梯度不是它这一步实际算出来的梯度，但梯度下降本身就是一个求近似解的过程，所以一般来说最终都还是会收敛。对于过期梯度的问题，目前还没有很好的解决，但可以设置一些类似过期梯度的阈值，比如一个梯度如果使用的参数版本小于当前ps上梯度的参数版本一定数量时，就放弃使用该梯度等等手段。

异步更新每次更新参数的速度都要比同步更新快得多，因为异步更新时不需要等待最慢的那个，没有短板效应。同样的，更新速度越快并不意味着收敛的快，也不意味着收敛时的精确度越高。

以上就是关于数据并行时的梯度下降。

至于该选择同步更新还是异步更新，这可以根据实际情况来调节，对于一般的小任务，比如单块GPU在几个小时内就可以跑完的任务，干脆就直接单GPU跑就可以了，这样就不会有很大的通信开销，也不会有过期梯度的问题。一般新手跑的模型基本上一块显卡足矣。

如果数据量很大，模型也很大，一定要使用分布式，那么该如何选择呢？对于这个问题，我不会很装X的回答，这得好好思考，因为没啥好思考的，同步还是异步我在上面都说明了，各有各的好处，同步每一次更新都会变慢，但总的更新次数可能会少，异步每次更新快，但总的更新次数可能会比较多，所以最终哪个收敛更快，这都不好说。也有种选择方式是，比如你有10台worker，每台worker上有4个GPU，那么你可以设置成每台worker内部的4块显卡同步，而worker与worker之间异步。

要相信不管你是用同步还是异步最后应该都会收敛，至于哪种方式更快，收敛时的准确率更高，这得去做尝试，不去做一遍谁都不知道（好像说的是废话）。

参考文献：

http://engineering.skymind.io/distributed-deep-learning-part-1-an-introduction-to-distributed-training-of-neural-networks