Caffe 多GPU卡相关代码理解

Caffe的NVIDIA GPU多卡训练使用NVIDIA nccl进行管理的，具体细节需要进一步深入的学习理解。当前只描述它的基本实现过程。

1. 1. 1. Class的基本关联关系图

1. 1. 1. 基本流程
2. 初始化阶段---构造NCCL

1. 初始化solver_。
2. 初始化size_(所有要学习的blobs的大小之和)。
3. 获取并指定GPU卡(默认为0)，分配size_大小的GPU内存data_(用于保存weight和bias的数据)，并把所有要学习的blobs的cpu_data拷贝到data_内。
4. 分配size_大小的GPU内存diff_(用于保存weight和bias的导数)，初始化为0；
5. 调用cudaStreamCreateWithFlags，创建cudaStream;
6. 调用apply_buffers，为blobs的每个blob的data和diff设置内存指针。

 

1. RUN阶段---NCCL::Run

为每个GPU卡创建并初始化Worker。具体处理如下：

1. 基于solver_、当前GPU的ID等信息，构造Worker。
2. 调用Worker的StartInternalThread操作，启动线程，该线程永久存在，除非运行结束或中断信号停止。该线程运行的主程序为Worker::InternalThreadEntry。该主程序Worker::InternalThreadEntry类似于Solver::Solve函数，唯一的不同是增加了通过NVIDIA提供的nccl库API接口对多卡的同步和callback处理。该主程序具体处理如下：

1. 指定该线程使用的GPU卡
2. 基于输入的solver_，创建该线程使用的Solver。
3. 基于Solver，构造该线程的NCCL，类似于初始化阶段的NCCL的操作。
4. Solver调用add_callback操作，这样Solver.callback_[i]->on_starth和Solver.callback_[i]->on_gradients_ready分别对应于NCCL::on_start和NCCL::on_gradients_ready操作；
5. Solver关联的Net调用add_after_backward操作，这样Net的after_backward_[i]->run就对应了NCCL::run操作；
6. 调用Solver的Step操作，进行训练或者test。

1. 等待中断信号退出，或运行结束退出。

因此，多卡训练密切相关的是NCCL::on_start、NCCL::on_gradients_ready操作和NCCL::run这三个操作。

1. NCCL::run是在每层进行backward结束后调用；对所有GPU卡的当前layer的所有blobs的diff求平均。
2. NCCL::on_start是空函数，无操作；
3. NCCL::on_gradients_ready是Solve在每次迭代的backward结束，更新参数之前调用。因为我们每层都做了NCCL::run，因此，不做此操作。