diff --git a/source/train/Data.py b/source/train/Data.py index db151d79cb..f3ecaeeb62 100644 --- a/source/train/Data.py +++ b/source/train/Data.py @@ -144,6 +144,12 @@ def get_test (self, ntests = -1) : self.modifier.modify_data(ret) return ret + def get_ntypes(self) : + if self.type_map is not None: + return len(self.type_map) + else: + return max(self.get_atom_type()) + 1 + def get_type_map(self) : return self.type_map diff --git a/source/train/DataSystem.py b/source/train/DataSystem.py index aca8994535..61d59f5cea 100644 --- a/source/train/DataSystem.py +++ b/source/train/DataSystem.py @@ -52,7 +52,7 @@ def __init__ (self, # natoms, nbatches ntypes = [] for ii in self.data_systems : - ntypes.append(np.max(ii.get_atom_type()) + 1) + ntypes.append(ii.get_ntypes()) self.sys_ntypes = max(ntypes) self.natoms = [] self.natoms_vec = [] diff --git a/source/train/Fitting.py b/source/train/Fitting.py index 5daaf386c3..960ff64e12 100644 --- a/source/train/Fitting.py +++ b/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]]) diff --git a/source/train/Network.py b/source/train/Network.py index fb87075942..ed188c085d 100644 --- a/source/train/Network.py +++ b/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