Update DescrptSeA.py

Author: GeiduanLiu

source/train/DescrptSeA.py

@@ -1,6 +1,6 @@
 import numpy as np
 from deepmd.env import tf
-from deepmd.common import ClassArg, get_activation_func
+from deepmd.common import ClassArg, get_activation_func， get_precision_func
 from deepmd.RunOptions import global_tf_float_precision
 from deepmd.RunOptions import global_np_float_precision
 from deepmd.env import op_module
@@ -19,7 +19,8 @@ def __init__ (self, jdata):
                .add('seed',     int) \
                .add('exclude_types', list, default = []) \
                .add('set_davg_zero', bool, default = False) \
-               .add('activation_function', str,    default = 'tanh')
+               .add('activation_function', str,    default = 'tanh') \
+               .add('precision', int, default = 0)
         class_data = args.parse(jdata)
         self.sel_a = class_data['sel']
         self.rcut_r = class_data['rcut']
@@ -30,6 +31,7 @@ def __init__ (self, jdata):
         self.seed = class_data['seed']
         self.trainable = class_data['trainable']
         self.filter_activation_fn = get_activation_func(class_data['activation_function'])
+        self.filter_precision = get_precision_func(class_data['precision'])
         exclude_types = class_data['exclude_types']
         self.exclude_types = set()
         for tt in exclude_types:
@@ -322,6 +324,7 @@ def _filter(self,
                    seed=None, 
                 trainable = True):
         # natom x (nei x 4)
+        inputs = tf.cast(inputs, self.filter_precision)
         shape = inputs.get_shape().as_list()
         outputs_size = [1] + self.filter_neuron
         outputs_size_2 = self.n_axis_neuron
@@ -343,18 +346,18 @@ def _filter(self,
               for ii in range(1, len(outputs_size)):
                 w = tf.get_variable('matrix_'+str(ii)+'_'+str(type_i), 
                                   [outputs_size[ii - 1], outputs_size[ii]], 
-                                  global_tf_float_precision,
+                                  self.filter_precision,
                                     tf.random_normal_initializer(stddev=stddev/np.sqrt(outputs_size[ii]+outputs_size[ii-1]), seed = seed), 
                                     trainable = trainable)
                 b = tf.get_variable('bias_'+str(ii)+'_'+str(type_i), 
                                   [1, outputs_size[ii]], 
-                                  global_tf_float_precision,
+                                  self.filter_precision,
                                   tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed), 
                                     trainable = trainable)
                 if self.filter_resnet_dt :
                     idt = tf.get_variable('idt_'+str(ii)+'_'+str(type_i), 
                                           [1, outputs_size[ii]], 
-                                          global_tf_float_precision,
+                                          self.filter_precision,
                                           tf.random_normal_initializer(stddev=0.001, mean = 1.0, seed = seed), 
                                           trainable = trainable)
                 if outputs_size[ii] == outputs_size[ii-1]:
@@ -409,6 +412,7 @@ def _filter_type_ext(self,
                            seed=None,
                          trainable = True):
         # natom x (nei x 4)
+        inputs = tf.cast(inputs, self.filter_precision)
         outputs_size = [1] + self.filter_neuron
         outputs_size_2 = self.n_axis_neuron
         with tf.variable_scope(name, reuse=reuse):
@@ -430,18 +434,18 @@ def _filter_type_ext(self,
             for ii in range(1, len(outputs_size)):
               w = tf.get_variable('matrix_'+str(ii)+'_'+str(type_i), 
                                 [outputs_size[ii - 1], outputs_size[ii]], 
-                                global_tf_float_precision,
+                                self.filter_precision,
                                   tf.random_normal_initializer(stddev=stddev/np.sqrt(outputs_size[ii]+outputs_size[ii-1]), seed = seed),
                                   trainable = trainable)
               b = tf.get_variable('bias_'+str(ii)+'_'+str(type_i), 
                                 [1, outputs_size[ii]], 
-                                global_tf_float_precision,
+                                self.filter_precision,
                                 tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed),
                                   trainable = trainable)
               if self.filter_resnet_dt :
                   idt = tf.get_variable('idt_'+str(ii)+'_'+str(type_i), 
                                         [1, outputs_size[ii]], 
-                                        global_tf_float_precision,
+                                        self.filter_precision,
                                         tf.random_normal_initializer(stddev=0.001, mean = 1.0, seed = seed),
                                         trainable = trainable)
               if outputs_size[ii] == outputs_size[ii-1]: