Merge pull request #295 from njzjz/master

merge master to devel
deepmodeling · Nov 24, 2020 · dc8469c · dc8469c
2 parents a84db43 + b64dbcf
commit dc8469c
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 4 deletions.
diff --git a/source/lib/include/common.h b/source/lib/include/common.h
@@ -1,16 +1,24 @@
 #pragma once
 
 #include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/version.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/shape_inference.h"
+#include <string>
 
 using namespace tensorflow;
 using namespace std;
 #include <tensorflow/core/graph/default_device.h>
 #include <tensorflow/core/graph/graph_def_builder.h>
 
+#if TF_MAJOR_VERSION >= 2 && TF_MINOR_VERSION >= 2
+typedef tensorflow::tstring STRINGTYPE;
+#else
+typedef std::string STRINGTYPE;
+#endif
+
 #include "NNPAtomMap.h"
 #include <vector>
 #include <string>

diff --git a/source/lib/src/DataModifier.cc b/source/lib/src/DataModifier.cc
@@ -37,7 +37,7 @@ init (const string & model,
   rcut = get_scalar<VALUETYPE>("descrpt_attr/rcut");
   cell_size = rcut;
   ntypes = get_scalar<int>("descrpt_attr/ntypes");
-  model_type = get_scalar<string>("model_attr/model_type");
+  model_type = get_scalar<STRINGTYPE>("model_attr/model_type");
   get_vector<int>(sel_type, "model_attr/sel_type");
   sort(sel_type.begin(), sel_type.end());
   inited = true;

diff --git a/source/lib/src/DeepTensor.cc b/source/lib/src/DeepTensor.cc
@@ -33,7 +33,7 @@ init (const string & model,
   rcut = get_scalar<VALUETYPE>("descrpt_attr/rcut");
   cell_size = rcut;
   ntypes = get_scalar<int>("descrpt_attr/ntypes");
-  model_type = get_scalar<string>("model_attr/model_type");
+  model_type = get_scalar<STRINGTYPE>("model_attr/model_type");
   odim = get_scalar<int>("model_attr/output_dim");
   get_vector<int>(sel_type, "model_attr/sel_type");
   inited = true;

diff --git a/source/lib/src/NNPInter.cc b/source/lib/src/NNPInter.cc
@@ -467,7 +467,7 @@ compute (ENERGYTYPE &			dener,
 void
 NNPInter::
 get_type_map(std::string & type_map){
-    type_map = get_scalar<std::string>("model_attr/tmap");
+    type_map = get_scalar<STRINGTYPE>("model_attr/tmap");
 }
 
 

diff --git a/source/train/Data.py b/source/train/Data.py
@@ -476,7 +476,7 @@ def load_energy(self,
         # ignore energy_file
         if coeff_atom_ener == 1:
             ener = np.sum(atom_ener, axis = 1)
-            coeff_atom_ener = 1
+            coeff_ener = 1
         # load energy_file
         else:
             coeff_ener, ener = self.load_data(set_name, energy_file, [nframes], False)