merge duplicated NeighborStat.get_stat (#1103)

Note: this is a simple fix to resolve #1088, but I think we should design a clear architecture to call neighbor stat.

This should reduce the half of the time, but it may be still too long. We can consider some better algorithm to calculate neighbour stat (like KDtree?) for further optimization.

Author: njzjz

deepmd/entrypoints/train.py

@@ -10,7 +10,7 @@
 from typing import Dict, List, Optional, Any
 
 from deepmd.common import data_requirement, expand_sys_str, j_loader, j_must_have
-from deepmd.env import tf, reset_default_tf_session_config
+from deepmd.env import tf, reset_default_tf_session_config, GLOBAL_TF_FLOAT_PRECISION
 from deepmd.infer.data_modifier import DipoleChargeModifier
 from deepmd.train.run_options import BUILD, CITATION, WELCOME, RunOptions
 from deepmd.train.trainer import DPTrainer
@@ -262,6 +262,16 @@ def get_nbor_stat(jdata, rcut):
     neistat = NeighborStat(ntypes, rcut)
 
     min_nbor_dist, max_nbor_size = neistat.get_stat(train_data)
+
+    # moved from traier.py as duplicated
+    # TODO: this is a simple fix but we should have a clear
+    #       architecture to call neighbor stat
+    tf.constant(min_nbor_dist,
+        name = 'train_attr/min_nbor_dist',
+        dtype = GLOBAL_TF_FLOAT_PRECISION)
+    tf.constant(max_nbor_size,
+        name = 'train_attr/max_nbor_size',
+        dtype = GLOBAL_TF_FLOAT_PRECISION)
     return min_nbor_dist, max_nbor_size
 
 def get_sel(jdata, rcut):

deepmd/train/trainer.py

@@ -313,16 +313,9 @@ def build (self,
             if self.run_opt.init_mode == 'init_from_frz_model':
                 self._init_from_frz_model()
 
-            self.neighbor_stat \
-                = NeighborStat(self.ntypes, self.descrpt.get_rcut())
-            self.min_nbor_dist, self.max_nbor_size \
-                = self.neighbor_stat.get_stat(data)
-            tf.constant(self.min_nbor_dist,
-                    name = 'train_attr/min_nbor_dist',
-                    dtype = GLOBAL_TF_FLOAT_PRECISION)
-            tf.constant(self.max_nbor_size,
-                    name = 'train_attr/max_nbor_size',
-                    dtype = GLOBAL_TF_FLOAT_PRECISION)
+            # neighbor_stat is moved to train.py as duplicated
+            # TODO: this is a simple fix but we should have a clear
+            #       architecture to call neighbor stat
         else :
             self.descrpt.enable_compression(self.model_param['compress']["min_nbor_dist"], self.model_param['compress']['model_file'], self.model_param['compress']['table_config'][0], self.model_param['compress']['table_config'][1], self.model_param['compress']['table_config'][2], self.model_param['compress']['table_config'][3])
             self.fitting.init_variables(get_fitting_net_variables(self.model_param['compress']['model_file']))