From 47336d899c3d1b8b1cb555e9b5bf562e49a712b8 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 17 Sep 2019 15:36:27 +0800 Subject: [PATCH] change interface of constructor, support multi-gpu mode in ModelDevi --- source/lib/include/NNPInter.h | 8 +-- source/lib/src/NNPInter.cc | 127 ++++++++++++++++++++-------------- 2 files changed, 78 insertions(+), 57 deletions(-) diff --git a/source/lib/include/NNPInter.h b/source/lib/include/NNPInter.h index 3c086a4a38..05376496d0 100644 --- a/source/lib/include/NNPInter.h +++ b/source/lib/include/NNPInter.h @@ -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, @@ -108,8 +108,8 @@ class NNPInterModelDevi { public: NNPInterModelDevi () ; - NNPInterModelDevi (const vector & models); - void init (const vector & models); + NNPInterModelDevi (const vector & models, const int & gpu_rank = 0); + void init (const vector & models, const int & gpu_rank = 0); public: void compute (ENERGYTYPE & ener, vector & force, diff --git a/source/lib/src/NNPInter.cc b/source/lib/src/NNPInter.cc index bc536b9d49..cc0c2e7111 100644 --- a/source/lib/src/NNPInter.cc +++ b/source/lib/src/NNPInter.cc @@ -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); @@ -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 @@ -667,18 +660,45 @@ NNPInterModelDevi () } NNPInterModelDevi:: -NNPInterModelDevi (const vector & models) +NNPInterModelDevi (const vector & 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 & 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(); @@ -686,10 +706,10 @@ NNPInterModelDevi (const vector & models) ntypes = get_ntypes(); inited = true; } - +#else void NNPInterModelDevi:: -init (const vector & models) +init (const vector & models, const int & gpu_rank) { assert (!inited); numb_models = models.size(); @@ -708,6 +728,7 @@ init (const vector & models) ntypes = get_ntypes(); inited = true; } +#endif VALUETYPE NNPInterModelDevi::