def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
polar_hat = label_dict[self.label_name]
polar = model_dict[self.tensor_name]
l2_loss = tf.reduce_mean(tf.square(self.scale * (polar - polar_hat)),
name='l2_' + suffix)
if not self.atomic:
atom_norm = 1. / global_cvt_2_tf_float(natoms[0])
l2_loss = l2_loss * atom_norm
self.l2_l = l2_loss
more_loss = {}
return l2_loss, more_loss
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_l = l2_loss
self.l2_more = more_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
self.l2_l = l2_loss
self.l2_more = more_loss
return l2_loss, more_loss