void NCCL<Dtype>::Run(const vector<int>& gpus, const char* restore) {
boost::barrier barrier(static_cast<int>(gpus.size()));
vector<NCCL<Dtype>*> nccls(gpus.size());
// Create workers
vector<shared_ptr<Worker<Dtype> > > workers(gpus.size());
for (int i = 1; i < gpus.size(); ++i) {
CUDA_CHECK(cudaSetDevice(gpus[i]));
Caffe::set_solver_rank(i);
Worker<Dtype>* w = new Worker<Dtype>(solver_, gpus[i], &barrier,
&nccls, restore);
w->StartInternalThread();
workers[i].reset(w);
}
CUDA_CHECK(cudaSetDevice(gpus[0]));
Caffe::set_solver_rank(0);
barrier_ = &barrier;
solver_->add_callback(this);
if (solver_->param().layer_wise_reduce()) {
solver_->net()->add_after_backward(this);
}
nccls[0] = this;
// Wait for workers
barrier.wait();
// Init NCCL
InitSingleProcess(&nccls);
barrier.wait();
// Run first solver on current thread
Broadcast();
solver_->Solve();
barrier.wait(); // Hangs without it when running tests
// Wait for shutdown
for (int i = 1; i < gpus.size(); ++i) {
workers[i]->StopInternalThread();
}
}
