Skip to content

Commit

Permalink
change interface of constructor, support multi-gpu mode in ModelDevi
Browse files Browse the repository at this point in the history
  • Loading branch information
Han Wang committed Sep 17, 2019
1 parent 47f0ed6 commit 47336d8
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 57 deletions.
8 changes: 4 additions & 4 deletions source/lib/include/NNPInter.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,8 @@ class NNPInter
{
public:
NNPInter () ;
NNPInter (const string & model);
void init (const string & model, const int & gpu_rank);
NNPInter (const string & model, const int & gpu_rank = 0);
void init (const string & model, const int & gpu_rank = 0);
void print_summary(const string &pre) const;
public:
void compute (ENERGYTYPE & ener,
Expand Down Expand Up @@ -108,8 +108,8 @@ class NNPInterModelDevi
{
public:
NNPInterModelDevi () ;
NNPInterModelDevi (const vector<string> & models);
void init (const vector<string> & models);
NNPInterModelDevi (const vector<string> & models, const int & gpu_rank = 0);
void init (const vector<string> & models, const int & gpu_rank = 0);
public:
void compute (ENERGYTYPE & ener,
vector<VALUETYPE> & force,
Expand Down
127 changes: 74 additions & 53 deletions source/lib/src/NNPInter.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -446,9 +446,50 @@ NNPInter ()
}

NNPInter::
NNPInter (const string & model)
NNPInter (const string & model, const int & gpu_rank)
: inited (false)
{
get_env_nthreads(num_intra_nthreads, num_inter_nthreads);
init(model, gpu_rank);
}

#ifdef USE_CUDA_TOOLKIT
void
NNPInter::
init (const string & model, const int & gpu_rank)
{
assert (!inited);
SessionOptions options;
options.config.set_inter_op_parallelism_threads(num_inter_nthreads);
options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
options.config.set_allow_soft_placement(true);
options.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction(0.9);
options.config.mutable_gpu_options()->set_allow_growth(true);

checkStatus (ReadBinaryProto(Env::Default(), model, &graph_def));
int gpu_num = -1;
cudaGetDeviceCount(&gpu_num);
// std::cout << "current number of devices: " << gpu_num << std::endl;
// set device to GPU only when at least GPU is found
if (gpu_num > 0) {
std::string str = "/gpu:";
str += std::to_string(gpu_rank % gpu_num);
graph::SetDefaultDevice(str, &graph_def);
// std::cout << "current device rank: " << str << std::endl;
}
checkStatus (NewSession(options, &session));
checkStatus (session->Create(graph_def));
rcut = get_rcut();
cell_size = rcut;
ntypes = get_ntypes();
inited = true;
}
#else
void
NNPInter::
init (const string & model, const int & gpu_rank)
{
assert (!inited);
SessionOptions options;
options.config.set_inter_op_parallelism_threads(num_inter_nthreads);
options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
Expand All @@ -460,54 +501,6 @@ NNPInter (const string & model)
ntypes = get_ntypes();
inited = true;
}

#ifdef USE_CUDA_TOOLKIT
void
NNPInter::
init (const string & model, const int & gpu_rank)
{
assert (!inited);
SessionOptions options;
options.config.set_inter_op_parallelism_threads(num_inter_nthreads);
options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
options.config.set_allow_soft_placement(true);
options.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction(0.9);
options.config.mutable_gpu_options()->set_allow_growth(true);

checkStatus (ReadBinaryProto(Env::Default(), model, &graph_def));
int gpu_num = 4;
cudaGetDeviceCount(&gpu_num);
std::string str = "/gpu:";
str += std::to_string(gpu_rank % gpu_num);
graph::SetDefaultDevice(str, &graph_def);
std::cout << "current device rank: " << str << std::endl;

checkStatus (NewSession(options, &session));
checkStatus (session->Create(graph_def));
rcut = get_rcut();
cell_size = rcut;
ntypes = get_ntypes();
inited = true;
//cudaGetDeviceCount(&_gpu_rank);
std::cout << "current device number: " << gpu_num << std::endl;
}
#else
void
NNPInter::
init (const string & model, const int & gpu_rank)
{
assert (!inited);
SessionOptions options;
options.config.set_inter_op_parallelism_threads(num_inter_nthreads);
options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
checkStatus (NewSession(options, &session));
checkStatus (ReadBinaryProto(Env::Default(), model, &graph_def));
checkStatus (session->Create(graph_def));
rcut = get_rcut();
cell_size = rcut;
ntypes = get_ntypes();
inited = true;
}
#endif

void
Expand Down Expand Up @@ -667,29 +660,56 @@ NNPInterModelDevi ()
}

NNPInterModelDevi::
NNPInterModelDevi (const vector<string> & models)
NNPInterModelDevi (const vector<string> & models, const int & gpu_rank)
: inited (false),
numb_models (0)
{
get_env_nthreads(num_intra_nthreads, num_inter_nthreads);
init(models, gpu_rank)
}

#ifdef USE_CUDA_TOOLKIT
void
NNPInterModelDevi::
init (const vector<string> & models, const int & gpu_rank)
{
assert (!inited);
numb_models = models.size();
sessions.resize(numb_models);
graph_defs.resize(numb_models);
SessionOptions options;
options.config.set_inter_op_parallelism_threads(num_inter_nthreads);
options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
options.config.set_allow_soft_placement(true);
options.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction(0.9);
options.config.mutable_gpu_options()->set_allow_growth(true);

for (unsigned ii = 0; ii < numb_models; ++ii){
checkStatus (NewSession(options, &(sessions[ii])));
checkStatus (ReadBinaryProto(Env::Default(), models[ii], &graph_defs[ii]));
}
int gpu_num = -1;
cudaGetDeviceCount(&gpu_num);
// std::cout << "current number of devices: " << gpu_num << std::endl;
for (unsigned ii = 0; ii < numb_models; ++ii){
// set device to GPU only when at least GPU is found
if (gpu_num > 0) {
std::string str = "/gpu:";
str += std::to_string(gpu_rank % gpu_num);
graph::SetDefaultDevice(str, &graph_def[ii]);
// std::cout << "current device rank: " << str << std::endl;
}
checkStatus (NewSession(options, &(sessions[ii])));
checkStatus (sessions[ii]->Create(graph_defs[ii]));
}
rcut = get_rcut();
cell_size = rcut;
ntypes = get_ntypes();
inited = true;
}

#else
void
NNPInterModelDevi::
init (const vector<string> & models)
init (const vector<string> & models, const int & gpu_rank)
{
assert (!inited);
numb_models = models.size();
Expand All @@ -708,6 +728,7 @@ init (const vector<string> & models)
ntypes = get_ntypes();
inited = true;
}
#endif

VALUETYPE
NNPInterModelDevi::
Expand Down

0 comments on commit 47336d8

Please sign in to comment.