fix Ener Model Infer

deepmodeling · Apr 18, 2021 · f813c77 · f813c77
1 parent 13b8e6f
commit f813c77
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 7 deletions.
diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py
@@ -118,8 +118,11 @@ def __init__ (self,
         self.t_ntypes = paddle.to_tensor(self.ntypes, dtype = "int32")
         self.t_ndescrpt = paddle.to_tensor(self.ndescrpt, dtype = "int32")
         self.t_sel = paddle.to_tensor(self.sel_a, dtype = "int32")
-        self.t_avg = paddle.to_tensor(np.zeros([self.ntypes, self.ndescrpt]), dtype = GLOBAL_PD_FLOAT_PRECISION)
-        self.t_std = paddle.to_tensor(np.ones([self.ntypes, self.ndescrpt]), dtype = GLOBAL_PD_FLOAT_PRECISION)
+
+        t_avg = paddle.to_tensor(np.zeros([self.ntypes, self.ndescrpt]), dtype = GLOBAL_PD_FLOAT_PRECISION)
+        t_std = paddle.to_tensor(np.ones([self.ntypes, self.ndescrpt]), dtype = GLOBAL_PD_FLOAT_PRECISION)
+        self.register_buffer("t_avg", t_avg)
+        self.register_buffer("t_std", t_std)
 
     def get_rcut (self) -> float:
         """
@@ -222,9 +225,7 @@ def compute_input_stats (self,
         if not self.set_davg_zero:
             self.davg = np.array(all_davg)
         self.dstd = np.array(all_dstd)
-
-        np.save("tf", self.davg)
-
+
         self.t_avg = paddle.to_tensor(self.davg, dtype = GLOBAL_NP_FLOAT_PRECISION)
         self.t_std = paddle.to_tensor(self.dstd, dtype = GLOBAL_NP_FLOAT_PRECISION)
 

diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py
@@ -26,7 +26,8 @@ def __init__(
         load_prefix: str = "load",
         default_tf_graph: bool = False
     ):
-        ##### hard code, should use dy2stat, avoid to build model #######
+        ##### Hard code, will change to use dy2stat, avoid to build model #######
+        ##### Now use paddle.load temporarily#######
         with open("out.json", 'r') as load_f:
             jdata = json.load(load_f)
 

diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py
@@ -253,7 +253,6 @@ def _eval_inner(
         # evaluate
         eval_inputs = {}
         eval_inputs['coord'] = paddle.to_tensor(np.reshape(coords, [-1]), dtype=GLOBAL_PD_FLOAT_PRECISION)
-        print(eval_inputs['coord'])
         eval_inputs['type'] = paddle.to_tensor(np.tile(atom_types, [nframes, 1]).reshape([-1]), dtype="int32")
         eval_inputs['natoms_vec'] = paddle.to_tensor(natoms_vec, dtype="int32")
         eval_inputs['box'] = paddle.to_tensor(np.reshape(cells , [-1]), dtype=GLOBAL_PD_FLOAT_PRECISION)