def _normalize_3d(self, a):
na = tf.norm(a, axis=1)
na = tf.tile(tf.reshape(na, [-1, 1]), tf.constant([1, 3]))
b = tf.divide(a, na)
return b
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