parallel.hpp
#ifndef CAFFE_PARALLEL_HPP_
#define CAFFE_PARALLEL_HPP_
#ifdef USE_NCCL
#include <boost/thread.hpp>
#include <string>
#include <vector>
#include "caffe/blob.hpp"
#include "caffe/common.hpp"
#include "caffe/internal_thread.hpp"
#include "caffe/layer.hpp"
#include "caffe/proto/caffe.pb.h"
#include "caffe/solver.hpp"
#include "caffe/syncedmem.hpp"
#include "caffe/util/blocking_queue.hpp"
#include "caffe/util/nccl.hpp"
namespace caffe {
// Represents a net parameters. Once a net is created, its parameter buffers can
// be replaced by ones from Params, to allow parallelization. Params ensures
// parameters are allocated in one consecutive array.
template<typename Dtype>
class Params {
public:
explicit Params(shared_ptr<Solver<Dtype> > root_solver);
virtual ~Params() {
}
inline size_t size() const {
return size_;
}
inline Dtype* data() const {
return data_;
}
inline Dtype* diff() const {
return diff_;
}
protected:
const size_t size_; // Size of buffers
Dtype* data_; // Network parameters
Dtype* diff_; // Gradient
DISABLE_COPY_AND_ASSIGN(Params);
};
// Params stored in GPU memory.
template<typename Dtype>
class GPUParams : public Params<Dtype> {
public:
GPUParams(shared_ptr<Solver<Dtype> > root_solver, int device);
virtual ~GPUParams();
void Configure(Solver<Dtype>* solver) const;
protected:
using Params<Dtype>::size_;
using Params<Dtype>::data_;
using Params<Dtype>::diff_;
};
template<typename Dtype>
class NCCL : public GPUParams<Dtype>,
public Solver<Dtype>::Callback,
public Net<Dtype>::Callback {
public:
explicit NCCL(shared_ptr<Solver<Dtype> > solver);
NCCL(shared_ptr<Solver<Dtype> > solver, const string& uid);
~NCCL();
boost::barrier* barrier();
void set_barrier(boost::barrier* value);
static void InitSingleProcess(vector<NCCL<Dtype>*>* nccls);
static string new_uid();
void Broadcast();
void Run(const vector<int>& gpus, const char* restore);
protected:
void Init();
void on_start() {}
void run(int layer); // Net callback
void on_gradients_ready();
ncclComm_t comm_;
cudaStream_t stream_;
shared_ptr<Solver<Dtype> > solver_;
// Should not be necessary, https://github.com/NVIDIA/nccl/issues/37
boost::barrier* barrier_;
using Params<Dtype>::size_;
using Params<Dtype>::data_;
using Params<Dtype>::diff_;
};
} // namespace caffe
#endif // USE_NCCL
#endif // header
caffe
A layer factory that allows one to register layers. During runtime, registered layers can be called b...
Definition: blob.hpp:14