Update DescrptSeR.py

Author: GeiduanLiu

source/train/DescrptSeR.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
@@ -18,7 +18,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_r = class_data['sel']
         self.rcut = class_data['rcut']
@@ -27,7 +28,8 @@ def __init__ (self, jdata):
         self.filter_resnet_dt = class_data['resnet_dt']
         self.seed = class_data['seed']        
         self.trainable = class_data['trainable']
-        self.filter_activation_fn = get_activation_func(class_data["activation_function"])   
+        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:
@@ -206,7 +208,7 @@ def _pass_filter(self,
                                  [ 0, start_index*      self.ndescrpt],
                                  [-1, natoms[2+type_i]* self.ndescrpt] )
             inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
-            layer = self._filter_r(inputs_i, type_i, name='filter_type_'+str(type_i)+suffix, natoms=natoms, reuse=reuse, seed = self.seed, trainable = trainable, activation_fn = self.filter_activation_fn)
+            layer = self._filter_r(tf.cast(inputs_i, self.filter_precision), type_i, name='filter_type_'+str(type_i)+suffix, natoms=natoms, reuse=reuse, seed = self.seed, trainable = trainable, activation_fn = self.filter_activation_fn)
             layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[2+type_i] * self.get_dim_out()])
             output.append(layer)
             start_index += natoms[2+type_i]
@@ -288,18 +290,18 @@ def _filter_r(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]: