卷积神经网络训练过程中的SGD的并行化设计

前段时间一直在关注 CNN 的实现，查看了 caffe 的代码以及 convnet2 的代码。目前对单机多卡的内容比较感兴趣，因此特别关注 convnet2 关于 multi-GPU 的支持。

其中 cuda-convnet2 的项目地址发布在：Google Code：cuda-convnet2

关于 multi-GPU 的一篇比较重要的论文就是：One weird trick for parallelizing convolutional neural networks

本文也将针对这篇文章给出分析。
 

1、简介

介绍一种卷积神经网络训练过程中的SGD的并行化方法。

两个变种

模型并行： 不同的 workers 训练模型的不同 patrs，比较适合神经元活动比较丰富的计算。

数据并行： 不同的 workers 训练不同的数据案例，比较适合 weight 矩阵比较多的计算。

2. 观察

现代卷积神经网络主要由两种层构成，他们具有不一样的属性和性能：
1）卷积层，占据了90% ~ 95% 的计算量，5% 的参数，但是对结果具有很大的表达能力。
2）全连接层，占据了 5% ~ 10% 的计算量， 95% 的参数，但是对于结果具有相对较小的表达的能力。

综上：卷积层计算量大，所需参数系数 W 少，全连接层计算量小，所需参数系数 W 多。因此对于卷积层适合使用数据并行，对于全连接层适合使用模型并行。

3. 推荐的算法

前向传播

1)K workers 中的每一个 worker 都提供不同的 128 个 examples 的 data batch，也就是每一个 worker 的数据都是不一样的。
2)每一个 worker 都在其 data batch 上计算卷积层。每一个 worker 的卷积层是按照顺序执行的。
3)全连接层的计算，分为以下三种方式：
(a) 每一个 worker 将其最后阶段的卷积层 activities 传递给其他的 worker。这些 workers 将这 128K 个 examples 配置成一个大的 batch，然后在这个 batch 上计算全连接层。
(b) 第一个 worker 将其最后阶段的卷积层 activities 传递给其他 workers，这些 workers 计算 128 个 examples 配置成的 batch 并且开始反向传递。（与这一次计算并行的同时，第二个 worker 将卷积层的 activities 传递到所有的 workers 中，即实现了 activities 传递和计算之间的流水线）。
(c) 全部的 workers 传递 128/K 个卷积层的 activities 到其他的 workers，计算方式同(b)。

对于上述(a~c)三种不同的全连接层实现方式，进行如下的分析：

(a) 当 128K images 配置到每一个 worker 的时候，所有有用的 work 必须要暂停。另外大的 batches 消耗大量的显存，这对于 GPU 显存有限制的设备是不希望发生的。另一方面，大的 batches 有利于 GPU 性能的发挥。

(b) 所有的 workers 轮流将他们的 activities 传播到所有的 workers 上面。这里轮流执行得到的最重要的结果是可以将大部分传播通信时间隐藏（因为他可以在上一次计算的时候，与计算一起并行处理，确切是 K-1 次通信时间可以隐藏）。这样做的意义非常的重大，可以实现一部分流水，使得通信时间隐藏，达到了很好的并行效果。

(c) 与 (b) 的方案类似。他的一个优势就是通信与计算的比例是常数 K，对于 (a) 和 (b)，与 K 成比例。这是因为 (a) 和 (b) 方案经常会受到每一个 worker 的输出带宽的限制，方案 (c) 可以利用所有的 workers 来完成这一任务，对于大的 K 会有非常大的优势。

反向传播

1) workers 按照通常的方式在全连接层计算梯度。
2) 按照前向传播中不同的实现方案，这里会有对应的三种方案：
(a) 每一个 worker 都为整个 128K examples 的 batch 计算了 activities 的梯度。所以每一个 worker 必须将每一个 example 的梯度传递给前向传播生成这个 example 的 worker 上面。这之后，卷积层的反向传播按照常规方式获取。
(b) 每一个 worker 已经为 128 examples 的 batch 计算了 activities 的梯度，然后将这些梯度值传递给每一个与这个 batch 相关的 workers 上。（在传播的同时，可以进行下一个 batch 的计算）。在经过 K 次前向和反向的传播，所有的梯度值都将被传播到卷积层。
(c) 与 (b) 类似。每一个 worker 计算了 128 examples batch 的梯度，这 128-example batch 来自于每一个 worker 的 128/K-examples，所以为了正确的分配每一个梯度，需要安装相反的操作来进行。

Weight 权值矩阵同步

一旦反向传播完成之后，workers 就可以更新 weight 矩阵了。在卷积层，workers 也必须同步 weight 矩阵。最简单的方式可以按如下进行：
1) 每一个 worker 指定到梯度矩阵的 1/K 进行同步。
2) 每一个 worker 从其他的 worker 累计对应的 1/K 的梯度。
3)每一个 worker 广播这 1/K 的梯度。

对于卷积层来说，这个比较容易实现，因为它的 weights 系数比较少。