set trainable of fitting layers

Author: Han Wang

source/train/Fitting.py

@@ -23,7 +23,8 @@ def __init__ (self, jdata, descrpt):
                .add('seed',             int)               \
                .add('atom_ener',        list,   default = [])\
                .add("activation_function", str,    default = "tanh")\
-               .add("precision",           str, default = "default")
+               .add("precision",           str, default = "default")\
+               .add("trainable",        [list, bool], default = True)
         class_data = args.parse(jdata)
         self.numb_fparam = class_data['numb_fparam']
         self.numb_aparam = class_data['numb_aparam']
@@ -32,7 +33,11 @@ def __init__ (self, jdata, descrpt):
         self.rcond = class_data['rcond']
         self.seed = class_data['seed']
         self.fitting_activation_fn = get_activation_func(class_data["activation_function"])
-        self.fitting_precision = get_precision(class_data['precision'])        
+        self.fitting_precision = get_precision(class_data['precision'])
+        self.trainable = class_data['trainable']
+        if type(self.trainable) is bool:
+            self.trainable = [self.trainable] * (len(self.n_neuron)+1)
+        assert(len(self.trainable) == len(self.n_neuron) + 1), 'length of trainable should be that of n_neuron + 1'
         self.atom_ener = []
         for at, ae in enumerate(class_data['atom_ener']):
             if ae is not None:
@@ -205,10 +210,10 @@ def build (self,
 
             for ii in range(0,len(self.n_neuron)) :
                 if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
-                    layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision)
+                    layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, trainable = self.trainable[ii])
                 else :
-                    layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision)
-            final_layer = one_layer(layer, 1, activation_fn = None, bavg = type_bias_ae, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision)
+                    layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, trainable = self.trainable[ii])
+            final_layer = one_layer(layer, 1, activation_fn = None, bavg = type_bias_ae, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, trainable = self.trainable[-1])
 
             if type_i < len(self.atom_ener) and self.atom_ener[type_i] is not None:
                 inputs_zero = tf.zeros_like(inputs_i, dtype=global_tf_float_precision)
@@ -219,10 +224,10 @@ def build (self,
                     layer = tf.concat([layer, ext_aparam], axis = 1)
                 for ii in range(0,len(self.n_neuron)) :
                     if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
-                        layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=True, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision)
+                        layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=True, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, trainable = self.trainable[ii])
                     else :
-                        layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=True, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision)
-                zero_layer = one_layer(layer, 1, activation_fn = None, bavg = type_bias_ae, name='final_layer_type_'+str(type_i)+suffix, reuse=True, seed = self.seed, precision = self.fitting_precision)
+                        layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=True, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, trainable = self.trainable[ii])
+                zero_layer = one_layer(layer, 1, activation_fn = None, bavg = type_bias_ae, name='final_layer_type_'+str(type_i)+suffix, reuse=True, seed = self.seed, precision = self.fitting_precision, trainable = self.trainable[-1])
                 final_layer += self.atom_ener[type_i] - zero_layer
 
             final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i]])

source/train/Network.py

@@ -13,23 +13,27 @@ def one_layer(inputs,
               reuse=None,
               seed=None, 
               use_timestep = False, 
+              trainable = True,
               useBN = False):
     with tf.variable_scope(name, reuse=reuse):
         shape = inputs.get_shape().as_list()
         w = tf.get_variable('matrix', 
                             [shape[1], outputs_size], 
                             precision,
-                            tf.random_normal_initializer(stddev=stddev/np.sqrt(shape[1]+outputs_size), seed = seed))
+                            tf.random_normal_initializer(stddev=stddev/np.sqrt(shape[1]+outputs_size), seed = seed), 
+                            trainable = trainable)
         b = tf.get_variable('bias', 
                             [outputs_size], 
                             precision,
-                            tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed))
+                            tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed), 
+                            trainable = trainable)
         hidden = tf.matmul(inputs, w) + b
         if activation_fn != None and use_timestep :
             idt = tf.get_variable('idt',
                                   [outputs_size],
                                   precision,
-                                  tf.random_normal_initializer(stddev=0.001, mean = 0.1, seed = seed))
+                                  tf.random_normal_initializer(stddev=0.001, mean = 0.1, seed = seed), 
+                                  trainable = trainable)
         if activation_fn != None:
             if useBN:
                 None