对境准备:对于多个GPU而言,一台机器2GPU,参数交换的流程图:

        Caffe参数交换源码分析

参数交换从main()进入train()函数,在train函数中找到对应源码为:

1 . . . . . 
2 if (gpus.size() > 1) {
3     caffe::P2PSync<float> sync(solver, NULL, solver->param());
4     sync.run(gpus);
5   } else {
6     LOG(INFO) << "Starting Optimization";
7     solver->Solve();
8   }

因为GPU的个数>1,所以执行sync(solver, NULL, solver->param())和run()函数,首先会执行P2PSync类的构造函数,然后执行run()函数,run函数的代码如下:

 1 void P2PSync<Dtype>::run(const vector<int>& gpus) {
 2   vector<DevicePair> pairs;
 3   DevicePair::compute(gpus, &pairs);
 4   SolverParameter param(solver_->param());
 5   vector<shared_ptr<P2PSync<Dtype> > > syncs(gpus.size());
 6 
 7   // Build the GPU tree by finding the parent for each solver
 8   for (int attempts = 0; attempts < pairs.size(); ++attempts) {. . . . . . .
 9   }
10   for (int i = 1; i < syncs.size(); ++i) {
11 syncs[i]->StartInternalThread();
12   }
13   solver_->Solve();
14   for (int i = 1; i < syncs.size(); ++i) {
15     syncs[i]->StopInternalThread();
16   }
17 }

在run()函数中,首先会执行compute()函数,该函数的作用是产生GPU Pairs,GPU Pairs的含义是[parent:child],对于2个GPU而言,GPU Pairs为[-1:0],[0:1],默认根GPU的parent是其本身。然后通过一个for循环构建GPU树,对于2个GPU而言,GPU树如下图所示:

        Caffe参数交换源码分析

接下来调用一个for循环为每个GPU开启一个线程,值得注意的是for循环是从i=1开始的,即为每个子GPU单独开启一个线程(这里为GPU1开启一个线程),也就是调用StartInternalThread()函数,该函数的代码如下:

1 void InternalThread::StartInternalThread() {. . . . . 
2   try {
3     thread_.reset(new boost::thread(&InternalThread::entry, this, device, mode,
4           rand_seed, solver_count, root_solver));
5   }. . . . . . .
6 }

该函数接着会执行entry()函数,该函数代码如下:

1 void InternalThread::entry(int device, Caffe::Brew mode, int rand_seed,
2 int solver_count, bool root_solver) {
3 . . . . . .
4   InternalThreadEntry();
5 }

该函数又会去调用InternalThreadEntry()函数,该函数是正式进入迭代运算的入口,代码如下:

 1 void P2PSync<Dtype>::InternalThreadEntry() {
 2   Caffe::SetDevice(solver_->param().device_id());
 3   CHECK(Caffe::root_solver());
 4   Caffe::set_root_solver(false);
 5   // See if there is a defined seed and reset random state if so
 6   if (solver_->param().random_seed() >= 0) {
 7     Caffe::set_random_seed(
 8         solver_->param().random_seed() + solver_->param().device_id());
 9   }
10   solver_->Step(solver_->param().max_iter() - initial_iter_);
11 }

GPU1调用Step()函数,进入迭代过程,见如下源码:

 1 void Solver<Dtype>::Step(int iters) {
 2   . . . . . . . . . .
 3   while (iter_ < stop_iter) {
 4 . . . . . . . . . .
 5 for (int i = 0; i < callbacks_.size(); ++i) {
 6       0_[i]->on_start();
 7     }
 8     const bool display = param_.display() && iter_ % param_.display() == 0;
 9     net_->set_debug_info(display && param_.debug_info());
10     // accumulate the loss and gradient
11     Dtype loss = 0;
12     for (int i = 0; i < param_.iter_size(); ++i) {
13       loss += net_->ForwardBackward(bottom_vec);//计算loss,一次前后向
14     }
15     loss /= param_.iter_size();//loss归一化
16     . . . . . . . 
17     for (int i = 0; i < callbacks_.size(); ++i) {
18       callbacks_[i]->on_gradients_ready();
19     }
20     ApplyUpdate();
21     . . . . . . . . . .
22     ++iter_;
23   }
24 }

整个Step函数的运行如上所示,首先根GPU(GPU0)有整个网络的网络参数,callbacks_.size()指的是GPU树的parent的个数(在这里是1),on_start()函数的作用就是把根GPU(GPU0)的网络参数分发到每一个子GPU(GPU1),GPU1会先进入这个函数,on_start()函数的部分代码如下:

1 void P2PSync<Dtype>::on_start() {
2   . . . . . . .
3   // Wait for update from parent
4   if (parent_) {
5     P2PSync<Dtype> *parent = queue_.pop();//取队列中的第一个gpu节点为根gpu
6     CHECK(parent == parent_);
7   }
8   . . . . . .

当执行到queue_.pop()时,会调用blocking_queue.cpp的pop()方法,pop()方法的内容如下:

1 T BlockingQueue<T>::pop(const string& log_on_wait) {
2   boost::mutex::scoped_lock lock(sync_->mutex_);
3   while (queue_.empty()) {
4     if (!log_on_wait.empty()) {
5       LOG_EVERY_N(INFO, 1000)<< log_on_wait;
6     }
7     sync_->condition_.wait(lock);//如果queue_为空,就一直阻塞。
8   }

该方法内部有wait()函数,因为此时queue_为空,所以GPU1就会被堵塞,因为GPU0GPU1是两个线程并行运行,所以GPU0会执行run()函数中的下一步,也就是solver_->Solve()Solve()函数的代码如下:

1 void Solver<Dtype>::Solve(const char* resume_file) {
2   int start_iter = iter_;
3 . . . . .
4   //LOG(INFO) <<"This is the sign of the train begin?********Ni****Jian*********";  //test for nijian
5   Step(param_.max_iter() - iter_);  
6   . . . . .
7 }

Solve()函数会调用Step()函数进入迭代过程,当GPU0进入on_start()函数后,会把队列中的GPU0出队列,同时会激活被堵塞的GPU1,接下来的on_start()函数代码如下:

 1  . . . . .
 2 // Update children
 3   for (int i = children_.size() - 1; i >= 0; i--) {
 4     Dtype* src = data_;
 5     Dtype* dst = children_[i]->data_;
 6 #ifdef DEBUG
 7    . . . .
 8 #endif
 9     CUDA_CHECK(cudaMemcpyAsync(dst, src, size_ * sizeof(Dtype),
10         cudaMemcpyDeviceToDevice, cudaStreamDefault));//每个子GPU把信息传入到根GPU,异步操作
11     CUDA_CHECK(cudaStreamSynchronize(cudaStreamDefault));//根GPU把信息同步传到各个子GPU
12     children_[i]->queue_.push(this);
13   }
14 #endif
15 }

在该部分代码中,src指的是GPU0的data(网络参数),dst指的是GPU1的data(网络参数),通过调用cudaMemcpyAsync()函数来放置一个请求,表示在cudaStreamDefault流中执行一次内存复制操作,然后调用cudaStreamSynchronize()等待cudaStreamDefault流中的操作完成后实现流的同步。经过这两个函数后,GPU0完成了把网络参数分发给GPU1,然后children_[i]->queue_.push(this)被执行后,会调用block_queue.cpp文件中的push函数激活GPU0的子GPU,即GPU1,同时把GPU1压入队列,此时队列中只有GPU1。

1 void BlockingQueue<T>::push(const T& t) {
2   boost::mutex::scoped_lock lock(sync_->mutex_);
3   queue_.push(t);
4   lock.unlock();
5   sync_->condition_.notify_one();
6 }

此时,多个GPU的参数分发过程已经完成,接下来GPU0和GPU1并行执行Step()函数的下一步,即:ForwardBackward(),该函数的代码如下:

1 Dtype ForwardBackward(const vector<Blob<Dtype>* > & bottom) {
2     Dtype loss;
3     Forward(bottom, &loss);
4     Backward();
5     return loss;
6   }

该函数的主要作用就是就是计算出loss和梯度diff,然后再接着执行Step()函数中的下一步,即:on_gradients_ready()函数,该函数分为两个部分,第一部分是多个GPU的梯度加和,第二部分是将计算后的梯度传给根GPU(GPU0)。第一部分的代码如下:

 1 void P2PSync<Dtype>::on_gradients_ready() {. . . . . . . .
 2   // Sum children gradients as they appear in the queue
 3   for (int i = 0; i < children_.size(); ++i) {
 4     P2PSync<Dtype> *child = queue_.pop();
 5     Dtype* src = child->parent_grads_;
 6     Dtype* dst = diff_;
 7 #ifdef DEBUG
 8     cudaPointerAttributes attributes;
 9     CUDA_CHECK(cudaPointerGetAttributes(&attributes, src));
10     CHECK(attributes.device == device);
11     CUDA_CHECK(cudaPointerGetAttributes(&attributes, dst));
12     CHECK(attributes.device == device);
13 #endif
14     caffe_gpu_add(size_, src, dst, dst);
15   }  

第一部分是多个GPU的梯度加和,因为GPU0和GPU1是并行计算的,如果GPU0执行到这里时,会使队列中仅有的GPU1出队列,然后通过调用caffe_gpu_add()函数,将一个GPU的梯度diff直接传给另一个GPU,不需要经过CPU通信,即GPU1把其计算的diff传给GPU0。如果是GPU1执行到这里时,因为GPU1没有子GPU,所以会直接跳过这一部分。第二部分的代码如下:

 1   if (parent_) {
 2     Dtype* src = diff_;
 3     Dtype* dst = parent_grads_;
 4 #ifdef DEBUG
 5 #endif
 6     CUDA_CHECK(cudaMemcpyAsync(dst, src, size_ * sizeof(Dtype),  //
 7         cudaMemcpyDeviceToDevice, cudaStreamDefault));
 8     CUDA_CHECK(cudaStreamSynchronize(cudaStreamDefault));
 9     parent_->queue_.push(this);
10   } else {
11     // Loss functions divide gradients by the batch size, so to compensate
12     // for split batch, the root solver divides by number of solvers.
13     caffe_gpu_scal(size_, Dtype(1.0 / Caffe::solver_count()), diff_);
14   }

如果是GPU0的话,会执行else,即caffe_gpu_scal(),该函数把得到的之前计算的梯度diff_和除以GPU的个数,来更新梯度。如果是GPU1的话,会执行if的语句,此时和on_start()函数分析类似,经过cudaMemcpyAsync()和cudaStreamSynchronize()函数操作之后,将GPU1中的梯度传送给GPU0,第二部分完成。

  接下来根GPU(GPU0)会得到所有的参数信息,会执行Step()函数的下一步,即执行ApplyUpdate()函数,该函数中有一个程序:CHECK(Caffe::root_solver()),会在根GPU中利用梯度下降法更新权重,计算参数,到此为止一次迭代完成,再进入下一次迭代时,根GPU已经保存了所有的网络参数,再继续迭代循环,直至结束。