NVcaffe源码阅读——Net&Solver

Solver

在caffe中，为了区分多GPU数据并行时，负责更新网络参数的solver被称作root solver。其他仅仅计算梯度用的solver为WorkerSolver，是Solver类的子类。nvcaffe将WorkerSolver类同Solver本身合并，将root solver和worker solver的概念融入到统一的框架下去，添加了is_root()等函数，也添加了Reduce()等用于并行处理的函数。

Step函数

nvcaffe中solver的Init()、InitTrainNet()和InitTestNet()没有明显改动。

1. 整理可学习参数的梯度空间

程序中使用了很多even(T val)函数。该函数返回一个val临近的偶数：

inline T even(T val) { return val & 1 ? val + 1 : val; }

对于Step()函数，nvcaffe首先调用了net.cpp所定义的InitializeLearnableDiffSpace()函数处理参数的梯度数值存储空间。和caffe使用ClearParamDiffs()函数对各个参数blob的diff空间进行写0覆盖相比，nvcaffe使用的InitializeLearnableDiffSpace()函数增加了内存对齐的操作（个人推断）。该函数是专门针对GPU显存的操作。在net.cpp的InitializeLearnableDiffSpace()函数中，nvcaffe使用了定义在gpu_memory.hpp中的Workspace结构体，将所有的参数梯度显存放在连续的空间之中。该函数首先计算每个待学习参数所占用的空间，计算方法是even(元素个数)*数据类型大小。个人推测之所以使用even()函数，是为了在使用FLOAT16类型时能够让显存对齐。之后使用这些计算好的大小来分配显存空间，最终仍然使用的是try_allocate()函数（详见”Blob的重新构建”的”Syncedmem”一节），并在此之后用0初始化空间。虽然各个梯度空间被像碎牛肉一样地压在了一起，但是各个空间的初始位置使用learnable_params_ptrs_这个指针数组保存了起来。

相比之下，caffe将梯度和参数本身以blob为单位在一起存放。nvcaffe则是将learnable_params_ptrs_的各个元素赋给每个blob的diff指针，但空间是连续的。

2. 多GPU之间的模型分发

如果使用了多GPU并行，则需要将net分发到各个设备上。两个框架的不同点在于callback_->on_start函数上，该函数定义在parallel.cpp的P2PSync空间里。caffe使用了cudaMemcpyAsync函数来同步数据，nvcaffe使用的是ncclBcast（如果编译时开启了使用NCCL）并辅以一些同步用的函数。

3. 异步并行更新权重

nvcaffe在step()函数中另一个显著的不同点在于权重的更新。在caffe中，权重的更新在每次前向后向传播后进行，使用ApplyUpdate()函数，而nvcaffe则是在循环迭代之前开启了一个新线程，专门负责权重的更新，调用solver.cpp中的Reduce()函数，以及进一步的net.cpp中的ReduceAndUpdate()函数：

    reduce_thread_.reset(new boost::thread(&Solver::Reduce, this,      Caffe::current_device(), mode, random_seed, solver_count, root_solver));    while (iter_ < stop_iter) {      ...    //start iteration

3.1 异步模型

step()函数的异步更新权重功能基本依赖于nvcaffe对net.cpp的改进。net.cpp维护了一个异步队列，该队列存储的元素是各个层的id号码：

BlockingQueue<int> reduction_queue_;

这个异步队列存储了经过网络反向计算了梯度之后了的、需要被更新权重的层的id。一方面，step()线程调用net.cpp的BackwardFromToAu()函数计算反向过程。每当计算完一个层的数据，且该层的参数需要被更新时，便将该层的id放入reduction_queue_中。每当计算完一个batch的数据后，放入一个END_OF_ITERATION标识符。另一方面，Reduce()线程则轮询reduction_queue_中的元素。一旦发现队列中有待处理参数信息时便调用实例化的solver中的ApplyUpdate()函数（例如，sgd_solver.cpp中的实现）。由于reduction_queue_本身是加锁的队列，因此能够为这两个线程提供异步的一致性。

step()函数最后调用了Net::Finalize()函数，以确保step()函数线程在Reduce()线程之后结束。

3.2 ReduceAndUpdate()函数与buckets

当使用多GPU进行并行训练时，用户可以使用NetParameter中的reduce_buckets选项来优化权重更新过程。
该过程的处理是在ReduceAndUpdate()函数当中。简单来讲，reduce_buckets用来设置当reduction_queue_累计了多少个待处理参数个数的时候才调用一次权重更新函数，类似于批处理的概念。在caffe.proto的注释中，作者建议reduce_buckets的默认参数对于大部分网络来说是比较好的设置。

multi-gpu TestAll&Test

nvcaffe也支持训练中多GPU的测试。在Test(const int test_net_id, const int iters, bool use_multi_gpu)函数中，如果使用了多GPU并行，则程序多执行一步同步各GPU测试结果的操作。TestAll(const int iters, bool use_multi_gpu)调用Test函数。在solver的构造函数中，use_multi_gpu的取值取决于Caffe::solver_count() > 1的结果，即是否使用了多GPU。