NVCaffe 0.16.2 多 GPU 训练过程代码分析

NVIDA在Caffe的基础上对其进行了优化，这篇文章主要是针对其多 GPU 训练过程中参数更新方式及通讯方法进行相关代码的学习，如有不正确之处请指正。 
先放主要的参考文章 
1. NVCaffe github 主页 
2. 博主 @KFXW 之前写了NVcaffe源码阅读系列文章，给了我很大启发，非常感谢！！ 
3. 另一位博主 @沤江一流 对 （Caffe，LeNet）的训练过程作了非常详细的介绍，前后向传播，权值更新几篇文章中让我学到了很多知识，同样非常感谢！！ 
4. 还参考了网络上其他博主的文章，很抱歉没有记录下来，但在此谢谢各位博主！

好了，进入正题，首先从主函数开始。 
主函数main()

int main(int argc, char** argv) {
  // Run tool or show usage.
  caffe::GlobalInit(&argc, &argv);
  // 设置设备 
  vector<int> gpus;
  get_gpus(&gpus);
#ifndef CPU_ONLY
  if (gpus.size() > 0) {
    Caffe::SetDevice(gpus[0]);
  }
#endif
  if (argc == 2) {
      // 若训练 caffe 的命令行为 ./build/tools/caffe train
      // 则这里 g_brew_map 的 key 值为 argv[1]，也即是 'train'，则实际调用了 train() 
      return GetBrewFunction(caffe::string(argv[1]))();  // ------->
  } else {
    gflags::ShowUsageWithFlagsRestrict(argv[0], "tools/caffe");
  }
}

RegisterBrewFunction 宏在每一个实现主要功能的函数之后将这个函数的名字和其对应的函数指针添加到了 g_brew_map 中, 具体分别为 train()，test()，device_query()，time() 这四个函数。

#define RegisterBrewFunction(func) \
namespace { \
class __Registerer_##func { \
 public: /* NOLINT */ \
  __Registerer_##func() { \
    g_brew_map[#func] = &func; \
  } \
}; \
__Registerer_##func g_registerer_##func; \
}

GetBrewFunction() 函数返回 g_brew_map[name], 即返回需要实现功能的函数。

static BrewFunction GetBrewFunction(const caffe::string& name) {
  if (g_brew_map.count(name)) {
    return g_brew_map[name];
  }
}  

进入 train() 函数，首先是从 solver.prototxt 文件中读取训练模型参数，并设置 Caffe 的 mode（GPU 还是 CPU）以及设备 id[s]，该部分代码省略，主要分析使用 Solver 类完成整个训练的过程。

int train() {
  // 通过调用 SolverRegistry 类的静态成员函数 CreateSolver() 得到一个指向 Solver 的指针来构造 shared_ptr 类型的 solver。
  // 这里的 solver_param 就是网络的模型及求解文件 solver.prototxt, 当多个 GPU 时，这里创建的 Solver 为训练过程的 root_solver, device_id = 0 (GPU0)。
  shared_ptr<caffe::Solver> solver(caffe::SolverRegistry::CreateSolver(solver_param)); //-----> solver_factory.hpp CreateSolver()
  solver->SetActionFunction(signal_handler.GetActionFunction());
  // 多 GPU 训练，需要涉及到 GPU 间通信与计算的异步处理问题。
  if (gpus.size() > 1)   {
    caffe::P2PManager p2p_mgr(solver, gpus.size(), solver->param()); 
                         //-----> parallel.cpp   P2PManager::P2PManager()
    p2p_mgr.Run(gpus); //-------> parallel.cpp  P2PManager::Run(const vector<int>& gpus)
  }   else {  // gpus.size() <= 1)
    LOG(INFO) << "Starting Optimization";
    // 调用 Solver 的 Solve() 方法，开始优化。
    solver->Solve();   //-------> solver.cpp  Solver::Solve(const char* resume_file)
  }
  LOG(INFO) << "Optimization Done in " << Caffe::time_from_init();
  return 0;
}

solver_factory.hpp 创建 Solver。

  static Solver* CreateSolver(const SolverParameter& param, size_t rank = 0U,
      Solver* root_solver = NULL) {
    const string& type = param.type();
    CreatorRegistry& registry = Registry();
    CHECK_EQ(registry.count(type), 1) << "Unknown solver type: " << type
        << " (known types: " << SolverTypeListString() << ")";
    Solver* solver = registry[type](param, rank, root_solver);
    return solver;
  }

尽管 solver 是一个指向基类 Solver 类型对象的指针，但由于 C++ 多态的特性，solver 这个智能指针调用各个成员函数时会调用到各个子类的函数。 
由于 caffe.proto 文件中默认的优化方法为 SGD，所以会实例化一个 SGDSolver 的对象（sgd_solvers.hpp）, SGDSolver 类继承于 Solver 类。 
class SGDSolver : public Solver 
构造函数为：

  explicit SGDSolver(const SolverParameter& param,
      size_t rank = 0U, Solver *root_solver = NULL)
      : Solver(param, rank, root_solver) { PreSolve(); }

因此，需要先调用父类 Solver 的构造函数，而 Solver 类中包含 Net 类对象，而 Net 类对象又包含了 Layers 类对象和 Blob 类对象。最终整个初始化的工作大概是： 

新建一个 SGDSolver 对象 -> 调用 SGDSolver 类的构造函数 -> 调用 Solver 类的构造函数 -> 新建 Net 类实例 -> 调用 Net 类的构造函数 -> 新建各个 Layer 的实例 -> 调用各个 Layer 类的构造函数 -> 设置每个 Blob，也由此完成整个网络的初始化。

parallel.cpp P2PManager 构造函数 
注意： caffe.cpp 中 创建的 solver 即为 root_solver

P2PManager::P2PManager(shared_ptr<Solver> root_solver,
    int nranks, const SolverParameter& solver_param) :
      nranks_(nranks),
      syncs_(nranks),
      root_solver_(root_solver)

parallel.cpp Run() 函数

void P2PManager::Run(const vector<int>& gpus) {
  ......
  SolverParameter param = root_solver_->param();
  this->shared_ = make_shared<SharedScores<float>>(nranks_);
  for (int i = 0; i < gpus.size(); ++i) {
    param.set_device_id(gpus[i]);
    // 返回一个 P2PSync 类型的 shared_ptr 智能指针 syncs_[i]
    // 每个 GPU 对应一个 P2PSync, 用于多 GPU 间的 P2P 异步
    syncs_[i] = make_shared<P2PSync>(this, root_solver_, i, gpus.size(), param);
                                 // -------> parallel.cpp  P2PSync::P2PSync()
#ifndef CPU_ONLY
#ifdef USE_NCCL
    syncs_[i]->aux_ = &nccl_id_;
#else
    LOG(FATAL) << "Multi-GPU execution not available - rebuild with USE_NCCL";
#endif  // USE_NCCL
#endif  // CPU_ONLY
    syncs_[i]->shared_ = this->shared_;
  }
  LOG(INFO)<< "Starting Optimization";
  for (int i = 0; i < syncs_.size(); ++i) {
    // 开始内部线程
    syncs_[i]->StartInternalThread(true, static_cast<uint64_t>(param.random_seed())); // -------> internal_thread.cpp 
       // InternalThread::StartInternalThread(bool set_cpu_affinity, uint64_t random_seed)
  }
  for (int i = 0; i < syncs_.size(); ++i) {
    syncs_[i]->WaitAll();
  }
  ......
}

P2PSync 类继承自 Solver::Callback 和 InternalThread 

class P2PSync : public Solver::Callback, public InternalThread 
构造函数：

P2PSync::P2PSync(P2PManager* mgr, shared_ptr<Solver> root_solver,
    int rank, int nranks, const SolverParamete