def build (self, learning_rate, natoms, model_dict, label_dict, suffix): coord = model_dict['coord'] energy = model_dict['energy'] atom_ener = model_dict['atom_ener'] nframes = tf.shape(atom_ener)[0] natoms = tf.shape(atom_ener)[1] # build energy dipole atom_ener0 = atom_ener - tf.reshape(tf.tile(tf.reshape(energy/global_cvt_2_ener_float(natoms), [-1, 1]), [1, natoms]), [nframes, natoms]) coord = tf.reshape(coord, [nframes, natoms, 3]) atom_ener0 = tf.reshape(atom_ener0, [nframes, 1, natoms]) ener_dipole = tf.matmul(atom_ener0, coord) ener_dipole = tf.reshape(ener_dipole, [nframes, 3]) energy_hat = label_dict['energy'] ener_dipole_hat = label_dict['energy_dipole'] find_energy = label_dict['find_energy'] find_ener_dipole = label_dict['find_energy_dipole'] l2_ener_loss = tf.reduce_mean( tf.square(energy - energy_hat), name='l2_'+suffix) ener_dipole_reshape = tf.reshape(ener_dipole, [-1]) ener_dipole_hat_reshape = tf.reshape(ener_dipole_hat, [-1]) l2_ener_dipole_loss = tf.reduce_mean( tf.square(ener_dipole_reshape - ener_dipole_hat_reshape), name='l2_'+suffix) # atom_norm_ener = 1./ global_cvt_2_ener_float(natoms[0]) atom_norm_ener = 1./ global_cvt_2_ener_float(natoms) pref_e = global_cvt_2_ener_float(find_energy * (self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * learning_rate / self.starter_learning_rate) ) pref_ed = global_cvt_2_tf_float(find_ener_dipole * (self.limit_pref_ed + (self.start_pref_ed - self.limit_pref_ed) * learning_rate / self.starter_learning_rate) ) l2_loss = 0 more_loss = {} l2_loss += atom_norm_ener * (pref_e * l2_ener_loss) l2_loss += global_cvt_2_ener_float(pref_ed * l2_ener_dipole_loss) more_loss['l2_ener_loss'] = l2_ener_loss more_loss['l2_ener_dipole_loss'] = l2_ener_dipole_loss self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss)) self.l2_loss_ener_summary = tf.summary.scalar('l2_ener_loss', tf.sqrt(l2_ener_loss) / global_cvt_2_tf_float(natoms[0])) self.l2_ener_dipole_loss_summary = tf.summary.scalar('l2_ener_dipole_loss', tf.sqrt(l2_ener_dipole_loss)) self.l2_l = l2_loss self.l2_more = more_loss return l2_loss, more_loss
def variable_summaries(var: tf.Variable, name: str): """Attach a lot of summaries to a Tensor (for TensorBoard visualization). Parameters ---------- var : tf.Variable [description] name : str variable name """ with tf.name_scope(name): mean = tf.reduce_mean(var) tf.summary.scalar('mean', mean) with tf.name_scope('stddev'): stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean))) tf.summary.scalar('stddev', stddev) tf.summary.scalar('max', tf.reduce_max(var)) tf.summary.scalar('min', tf.reduce_min(var)) tf.summary.histogram('histogram', var)
def build(self, learning_rate, natoms, model_dict, label_dict, suffix): polar_hat = label_dict[self.label_name] atomic_polar_hat = label_dict["atomic_" + self.label_name] polar = tf.reshape(model_dict[self.tensor_name], [-1]) find_global = label_dict['find_' + self.label_name] find_atomic = label_dict['find_atomic_' + self.label_name] # YHT: added for global / local dipole combination l2_loss = global_cvt_2_tf_float(0.0) more_loss = { "local_loss": global_cvt_2_tf_float(0.0), "global_loss": global_cvt_2_tf_float(0.0) } if self.local_weight > 0.0: local_loss = global_cvt_2_tf_float(find_atomic) * tf.reduce_mean( tf.square(self.scale * (polar - atomic_polar_hat)), name='l2_' + suffix) more_loss['local_loss'] = local_loss l2_loss += self.local_weight * local_loss self.l2_loss_local_summary = tf.summary.scalar( 'l2_local_loss', tf.sqrt(more_loss['local_loss'])) if self.global_weight > 0.0: # Need global loss atoms = 0 if self.type_sel is not None: for w in self.type_sel: atoms += natoms[2 + w] else: atoms = natoms[0] nframes = tf.shape(polar)[0] // self.tensor_size // atoms # get global results global_polar = tf.reshape( tf.reduce_sum(tf.reshape(polar, [nframes, -1, self.tensor_size]), axis=1), [-1]) #if self.atomic: # If label is local, however # global_polar_hat = tf.reshape(tf.reduce_sum(tf.reshape( # polar_hat, [nframes, -1, self.tensor_size]), axis=1),[-1]) #else: # global_polar_hat = polar_hat global_loss = global_cvt_2_tf_float(find_global) * tf.reduce_mean( tf.square(self.scale * (global_polar - polar_hat)), name='l2_' + suffix) more_loss['global_loss'] = global_loss self.l2_loss_global_summary = tf.summary.scalar( 'l2_global_loss', tf.sqrt(more_loss['global_loss']) / global_cvt_2_tf_float(atoms)) # YWolfeee: should only consider atoms with dipole, i.e. atoms # atom_norm = 1./ global_cvt_2_tf_float(natoms[0]) atom_norm = 1. / global_cvt_2_tf_float(atoms) global_loss *= atom_norm l2_loss += self.global_weight * global_loss self.l2_more = more_loss self.l2_l = l2_loss self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss)) return l2_loss, more_loss
def build (self, learning_rate, natoms, model_dict, label_dict, suffix): energy = model_dict['energy'] force = model_dict['force'] virial = model_dict['virial'] atom_ener = model_dict['atom_ener'] energy_hat = label_dict['energy'] force_hat = label_dict['force'] virial_hat = label_dict['virial'] atom_ener_hat = label_dict['atom_ener'] atom_pref = label_dict['atom_pref'] find_energy = label_dict['find_energy'] find_force = label_dict['find_force'] find_virial = label_dict['find_virial'] find_atom_ener = label_dict['find_atom_ener'] find_atom_pref = label_dict['find_atom_pref'] l2_ener_loss = tf.reduce_mean( tf.square(energy - energy_hat), name='l2_'+suffix) force_reshape = tf.reshape (force, [-1]) force_hat_reshape = tf.reshape (force_hat, [-1]) atom_pref_reshape = tf.reshape (atom_pref, [-1]) diff_f = force_hat_reshape - force_reshape if self.relative_f is not None: force_hat_3 = tf.reshape(force_hat, [-1, 3]) norm_f = tf.reshape(tf.norm(force_hat_3, axis = 1), [-1, 1]) + self.relative_f diff_f_3 = tf.reshape(diff_f, [-1, 3]) diff_f_3 = diff_f_3 / norm_f diff_f = tf.reshape(diff_f_3, [-1]) l2_force_loss = tf.reduce_mean(tf.square(diff_f), name = "l2_force_" + suffix) l2_pref_force_loss = tf.reduce_mean(tf.multiply(tf.square(diff_f), atom_pref_reshape), name = "l2_pref_force_" + suffix) virial_reshape = tf.reshape (virial, [-1]) virial_hat_reshape = tf.reshape (virial_hat, [-1]) l2_virial_loss = tf.reduce_mean (tf.square(virial_hat_reshape - virial_reshape), name = "l2_virial_" + suffix) atom_ener_reshape = tf.reshape (atom_ener, [-1]) atom_ener_hat_reshape = tf.reshape (atom_ener_hat, [-1]) l2_atom_ener_loss = tf.reduce_mean (tf.square(atom_ener_hat_reshape - atom_ener_reshape), name = "l2_atom_ener_" + suffix) atom_norm = 1./ global_cvt_2_tf_float(natoms[0]) atom_norm_ener = 1./ global_cvt_2_ener_float(natoms[0]) pref_e = global_cvt_2_ener_float(find_energy * (self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * learning_rate / self.starter_learning_rate) ) pref_f = global_cvt_2_tf_float(find_force * (self.limit_pref_f + (self.start_pref_f - self.limit_pref_f) * learning_rate / self.starter_learning_rate) ) pref_v = global_cvt_2_tf_float(find_virial * (self.limit_pref_v + (self.start_pref_v - self.limit_pref_v) * learning_rate / self.starter_learning_rate) ) pref_ae= global_cvt_2_tf_float(find_atom_ener * (self.limit_pref_ae+ (self.start_pref_ae-self.limit_pref_ae) * learning_rate / self.starter_learning_rate) ) pref_pf= global_cvt_2_tf_float(find_atom_pref * (self.limit_pref_pf+ (self.start_pref_pf-self.limit_pref_pf) * learning_rate / self.starter_learning_rate) ) l2_loss = 0 more_loss = {} if self.has_e : l2_loss += atom_norm_ener * (pref_e * l2_ener_loss) more_loss['l2_ener_loss'] = l2_ener_loss if self.has_f : l2_loss += global_cvt_2_ener_float(pref_f * l2_force_loss) more_loss['l2_force_loss'] = l2_force_loss if self.has_v : l2_loss += global_cvt_2_ener_float(atom_norm * (pref_v * l2_virial_loss)) more_loss['l2_virial_loss'] = l2_virial_loss if self.has_ae : l2_loss += global_cvt_2_ener_float(pref_ae * l2_atom_ener_loss) more_loss['l2_atom_ener_loss'] = l2_atom_ener_loss if self.has_pf : l2_loss += global_cvt_2_ener_float(pref_pf * l2_pref_force_loss) more_loss['l2_pref_force_loss'] = l2_pref_force_loss # only used when tensorboard was set as true self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss)) self.l2_loss_ener_summary = tf.summary.scalar('l2_ener_loss', global_cvt_2_tf_float(tf.sqrt(l2_ener_loss)) / global_cvt_2_tf_float(natoms[0])) self.l2_loss_force_summary = tf.summary.scalar('l2_force_loss', tf.sqrt(l2_force_loss)) self.l2_loss_virial_summary = tf.summary.scalar('l2_virial_loss', tf.sqrt(l2_virial_loss) / global_cvt_2_tf_float(natoms[0])) self.l2_l = l2_loss self.l2_more = more_loss return l2_loss, more_loss