remove dependency of TF headers from C++ public headers

doc/inference/cxx.md

@@ -17,7 +17,7 @@ where `e`, `f` and `v` are predicted energy, force and virial of the system, res
 
 You can compile `infer_water.cpp` using `gcc`:
 ```sh
-gcc infer_water.cpp -D HIGH_PREC -L $deepmd_root/lib -L $tensorflow_root/lib -I $deepmd_root/include -I $tensorflow_root/include -Wl,--no-as-needed -ldeepmd_cc -lstdc++ -ltensorflow_cc -Wl,-rpath=$deepmd_root/lib -Wl,-rpath=$tensorflow_root/lib -o infer_water
+gcc infer_water.cpp -D HIGH_PREC -L $deepmd_root/lib -L $tensorflow_root/lib -I $deepmd_root/include -Wl,--no-as-needed -ldeepmd_cc -lstdc++ -ltensorflow_cc -Wl,-rpath=$deepmd_root/lib -Wl,-rpath=$tensorflow_root/lib -o infer_water
 ```
 and then run the program:
 ```sh

source/api_cc/CMakeLists.txt

@@ -17,14 +17,18 @@ add_library(${libname} SHARED ${LIB_SRC})
 
 # link: libdeepmd libdeepmd_op libtensorflow_cc libtensorflow_framework
 target_link_libraries (${libname} PUBLIC ${LIB_DEEPMD}	${TensorFlow_LIBRARY} ${TensorFlowFramework_LIBRARY})
-target_include_directories(${libname} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include ${CMAKE_CURRENT_BINARY_DIR} ${TensorFlow_INCLUDE_DIRS})
+target_include_directories(${libname} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include ${CMAKE_CURRENT_BINARY_DIR})
+target_include_directories(${libname} PRIVATE ${TensorFlow_INCLUDE_DIRS})
 
 set_target_properties(
   ${libname} 
   PROPERTIES 
-  COMPILE_DEFINITIONS ${prec_def}
   INSTALL_RPATH "$ORIGIN;${TensorFlow_LIBRARY_PATH}"
 )
+target_compile_definitions(${libname}
+  PUBLIC ${prec_def}
+  PRIVATE TF_PRIVATE
+)
 
 install(TARGETS ${libname} DESTINATION lib/)
 

source/api_cc/include/DataModifier.h

@@ -32,7 +32,7 @@ class DipoleChargeModifier
   tensorflow::Session* session;
   std::string name_scope, name_prefix;
   int num_intra_nthreads, num_inter_nthreads;
-  tensorflow::GraphDef graph_def;
+  tensorflow::GraphDef* graph_def;
   bool inited;
   VALUETYPE rcut;
   VALUETYPE cell_size;

source/api_cc/include/DeepPot.h

@@ -179,7 +179,7 @@ class DeepPot
 private:
   tensorflow::Session* session;
   int num_intra_nthreads, num_inter_nthreads;
-  tensorflow::GraphDef graph_def;
+  tensorflow::GraphDef* graph_def;
   bool inited;
   template<class VT> VT get_scalar(const std::string & name) const;
   // VALUETYPE get_rcut () const;
@@ -401,7 +401,7 @@ class DeepPotModelDevi
   unsigned numb_models;
   std::vector<tensorflow::Session*> sessions;
   int num_intra_nthreads, num_inter_nthreads;
-  std::vector<tensorflow::GraphDef> graph_defs;
+  std::vector<tensorflow::GraphDef*> graph_defs;
   bool inited;
   template<class VT> VT get_scalar(const std::string name) const;
   // VALUETYPE get_rcut () const;
@@ -426,7 +426,6 @@ class DeepPotModelDevi
 
   // function used for nborlist copy
   std::vector<std::vector<int> > get_sel() const;
-  void cum_sum(const std::vector<std::vector<tensorflow::int32> > n_sel);
 };
 }
 

source/api_cc/include/DeepTensor.h

@@ -160,7 +160,7 @@ class DeepTensor
   tensorflow::Session* session;
   std::string name_scope;
   int num_intra_nthreads, num_inter_nthreads;
-  tensorflow::GraphDef graph_def;
+  tensorflow::GraphDef* graph_def;
   bool inited;
   VALUETYPE rcut;
   VALUETYPE cell_size;

source/api_cc/include/common.h

@@ -8,21 +8,15 @@
 #include "AtomMap.h"
 #include "errors.h"
 
-#include "tensorflow/core/platform/env.h"
-#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/public/version.h"
-#include <tensorflow/core/graph/default_device.h>
-#include <tensorflow/core/graph/graph_def_builder.h>
+#ifdef TF_PRIVATE
+#include "tf_private.h"
+#else
+#include "tf_public.h"
+#endif
 
 
 namespace deepmd{
 
-#if TF_MAJOR_VERSION >= 2 && TF_MINOR_VERSION >= 2
-typedef tensorflow::tstring STRINGTYPE;
-#else
-typedef std::string STRINGTYPE;
-#endif
-
 #ifdef HIGH_PREC
 typedef double VALUETYPE;
 typedef double ENERGYTYPE;
@@ -188,5 +182,22 @@ session_input_tensors (std::vector<std::pair<std::string, tensorflow::Tensor>> &
 		       const int			nghost,
 		       const int			ago,
 		       const std::string		scope = "");
+
+/**
+* @brief Read model file to a string.
+* @param[in] model Path to the model.
+* @param[out] file_content Content of the model file.
+**/
+void
+read_file_to_string(std::string model, std::string & file_content);
+
+
+/**
+* @brief Convert pbtxt to pb.
+* @param[in] fn_pb_txt Filename of the pb txt file.
+* @param[out] fn_pb Filename of the pb file.
+**/
+void
+convert_pbtxt_to_pb(std::string fn_pb_txt, std::string fn_pb);
 }
 

source/api_cc/include/tf_private.h

@@ -0,0 +1,19 @@
+/**
+ * @file tf_private.h
+ * @brief This file includes TensorFlow headers used for compilation.
+ * 
+ */
+
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/version.h"
+#include <tensorflow/core/graph/default_device.h>
+#include <tensorflow/core/graph/graph_def_builder.h>
+
+namespace deepmd {
+#if TF_MAJOR_VERSION >= 2 && TF_MINOR_VERSION >= 2
+typedef tensorflow::tstring STRINGTYPE;
+#else
+typedef std::string STRINGTYPE;
+#endif
+}

source/api_cc/include/tf_public.h

@@ -0,0 +1,15 @@
+/**
+ * @file tf_public.h
+ * @brief This file declares incompleted TensorFlow class used for public headers.
+ * 
+ */
+
+// skip if TF headers have been included
+#ifndef TF_MAJOR_VERSION
+namespace tensorflow{
+    class Session;
+    class Tensor;
+    class GraphDef;
+    class Status;
+}
+#endif

source/api_cc/src/DataModifier.cc

@@ -5,15 +5,17 @@ using namespace tensorflow;
 
 DipoleChargeModifier::
 DipoleChargeModifier()
-    : inited (false)
+    : inited (false),
+      graph_def(new GraphDef())
 {
 }
 
 DipoleChargeModifier::
 DipoleChargeModifier(const std::string & model, 
 	     const int & gpu_rank, 
 	     const std::string &name_scope_)
-    : inited (false), name_scope(name_scope_)
+    : inited (false), name_scope(name_scope_),
+      graph_def(new GraphDef())
 {
   init(model, gpu_rank);  
 }
@@ -35,8 +37,8 @@ init (const std::string & model,
   options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
   deepmd::load_op_library();
   deepmd::check_status(NewSession(options, &session));
-  deepmd::check_status(ReadBinaryProto(Env::Default(), model, &graph_def));
-  deepmd::check_status(session->Create(graph_def));  
+  deepmd::check_status(ReadBinaryProto(Env::Default(), model, graph_def));
+  deepmd::check_status(session->Create(*graph_def));  
   // int nnodes = graph_def.node_size();
   // for (int ii = 0; ii < nnodes; ++ii){
   //   cout << ii << " \t " << graph_def.node(ii).name() << endl;

source/api_cc/src/DeepPot.cc

@@ -162,13 +162,15 @@ static void run_model (ENERGYTYPE   &		dener,
 
 DeepPot::
 DeepPot ()
-    : inited (false), init_nbor (false)
+    : inited (false), init_nbor (false),
+      graph_def(new GraphDef())
 {
 }
 
 DeepPot::
 DeepPot (const std::string & model, const int & gpu_rank, const std::string & file_content)
-    : inited (false), init_nbor (false)
+    : inited (false), init_nbor (false),
+      graph_def(new GraphDef())
 {
   init(model, gpu_rank, file_content);  
 }
@@ -190,9 +192,9 @@ init (const std::string & model, const int & gpu_rank, const std::string & file_
   deepmd::load_op_library();
 
   if(file_content.size() == 0)
-    check_status (ReadBinaryProto(Env::Default(), model, &graph_def));
+    check_status (ReadBinaryProto(Env::Default(), model, graph_def));
   else
-    graph_def.ParseFromString(file_content);
+    (*graph_def).ParseFromString(file_content);
   int gpu_num = -1;
   #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
   DPGetDeviceCount(gpu_num); // check current device environment
@@ -203,11 +205,11 @@ init (const std::string & model, const int & gpu_rank, const std::string & file_
     DPErrcheck(DPSetDevice(gpu_rank % gpu_num));
     std::string str = "/gpu:";
     str += std::to_string(gpu_rank % gpu_num);
-    graph::SetDefaultDevice(str, &graph_def);
+    graph::SetDefaultDevice(str, graph_def);
   }
   #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
   check_status (NewSession(options, &session));
-  check_status (session->Create(graph_def));
+  check_status (session->Create(*graph_def));
   rcut = get_scalar<VALUETYPE>("descrpt_attr/rcut");
   cell_size = rcut;
   ntypes = get_scalar<int>("descrpt_attr/ntypes");
@@ -282,7 +284,7 @@ std::string graph_info(const GraphDef & graph_def) {
 // init the tmp array data
 std::vector<int> DeepPot::get_sel_a () const {
     std::vector<int> sel_a;
-    std::istringstream is(graph_info(graph_def));
+    std::istringstream is(graph_info(*graph_def));
     std::string line = "";
     while(is >> line) {
         if (line.find("sel_a") != line.npos) {
@@ -548,10 +550,11 @@ init (const std::vector<std::string> & models, const int & gpu_rank, const std::
   options.config.set_inter_op_parallelism_threads(num_inter_nthreads);
   options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
   for (unsigned ii = 0; ii < numb_models; ++ii){
+    graph_defs[ii] = new GraphDef();
     if (file_contents.size() == 0)
-      check_status (ReadBinaryProto(Env::Default(), models[ii], &graph_defs[ii]));
+      check_status (ReadBinaryProto(Env::Default(), models[ii], graph_defs[ii]));
     else
-      graph_defs[ii].ParseFromString(file_contents[ii]);
+      (*graph_defs[ii]).ParseFromString(file_contents[ii]);
   }
   #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
   if (gpu_num > 0) {
@@ -566,10 +569,10 @@ init (const std::vector<std::string> & models, const int & gpu_rank, const std::
     if (gpu_num > 0) {
       std::string str = "/gpu:";
       str += std::to_string(gpu_rank % gpu_num);
-      graph::SetDefaultDevice(str, &graph_defs[ii]);
+      graph::SetDefaultDevice(str, &(*graph_defs[ii]));
     }
     check_status (NewSession(options, &(sessions[ii])));
-    check_status (sessions[ii]->Create(graph_defs[ii]));
+    check_status (sessions[ii]->Create(*graph_defs[ii]));
   }
   rcut = get_scalar<VALUETYPE>("descrpt_attr/rcut");
   cell_size = rcut;
@@ -620,7 +623,7 @@ get_sel () const
     std::vector<std::vector<int> > sec;
     for (int ii = 0; ii < numb_models; ii++) {
         std::vector<int> sel;
-        std::istringstream is(graph_info(graph_defs[ii]));
+        std::istringstream is(graph_info(*graph_defs[ii]));
         std::string line = "";
         while(is >> line) {
             if (line.find("sel") != line.npos) {
@@ -643,20 +646,6 @@ get_sel () const
     return sec;
 }
 
-void  
-DeepPotModelDevi::
-cum_sum (const std::vector<std::vector<int32> > n_sel) 
-{
-    for (int ii = 0; ii < numb_models; ++ii) {
-        std::vector<int> _sec;
-        _sec.resize (n_sel[ii].size() + 1);
-        _sec[0] = 0;
-        for (int jj = 1; jj < _sec.size(); ++jj) {
-            _sec[jj] = _sec[jj-1] + n_sel[ii][jj-1];
-        }
-        sec.push_back(_sec);
-    }
-}
 
 void
 DeepPotModelDevi::

source/api_cc/src/DeepTensor.cc

@@ -5,15 +5,17 @@ using namespace tensorflow;
 
 DeepTensor::
 DeepTensor()
-    : inited (false)
+    : inited (false),
+      graph_def(new GraphDef())
 {
 }
 
 DeepTensor::
 DeepTensor(const std::string & model, 
 	   const int & gpu_rank, 
 	   const std::string &name_scope_)
-    : inited (false), name_scope(name_scope_)
+    : inited (false), name_scope(name_scope_),
+      graph_def(new GraphDef())
 {
   init(model, gpu_rank);  
 }
@@ -35,8 +37,8 @@ init (const std::string & model,
   options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
   deepmd::load_op_library();
   deepmd::check_status (NewSession(options, &session));
-  deepmd::check_status (ReadBinaryProto(Env::Default(), model, &graph_def));
-  deepmd::check_status (session->Create(graph_def));  
+  deepmd::check_status (ReadBinaryProto(Env::Default(), model, graph_def));
+  deepmd::check_status (session->Create(*graph_def));  
   rcut = get_scalar<VALUETYPE>("descrpt_attr/rcut");
   cell_size = rcut;
   ntypes = get_scalar<int>("descrpt_attr/ntypes");

source/api_cc/src/common.cc

@@ -2,6 +2,9 @@
 #include "AtomMap.h"
 #include "device.h"
 #include <dlfcn.h>
+#include <fcntl.h>
+#include "google/protobuf/text_format.h"
+#include "google/protobuf/io/zero_copy_stream_impl.h"
 
 using namespace tensorflow;
 
@@ -855,3 +858,25 @@ select_map_inv<deepmd::STRINGTYPE >(
     const typename std::vector<deepmd::STRINGTYPE >::const_iterator in, 
     const std::vector<int > & idx_map, 
     const int & stride);
+
+
+void
+deepmd::
+read_file_to_string(std::string model, std::string & file_content)
+{
+  deepmd::check_status(tensorflow::ReadFileToString(tensorflow::Env::Default(), model, &file_content));
+}
+
+
+void
+deepmd::
+convert_pbtxt_to_pb(std::string fn_pb_txt, std::string fn_pb)
+{
+    int fd = open(fn_pb_txt.c_str(), O_RDONLY);
+    tensorflow::protobuf::io::ZeroCopyInputStream* input = new tensorflow::protobuf::io::FileInputStream(fd);
+    tensorflow::GraphDef graph_def;
+    tensorflow::protobuf::TextFormat::Parse(input, &graph_def);
+    delete input;
+    std::fstream output(fn_pb, std::ios::out | std::ios::trunc | std::ios::binary);
+    graph_def.SerializeToOstream(&output);
+}