diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index bf88a9a1e6..e5d291e7dc 100644 --- a/deepmd/entrypoints/main.py +++ b/deepmd/entrypoints/main.py @@ -232,10 +232,10 @@ def parse_args(args: Optional[List[str]] = None): ) parser_tst.add_argument( "-a", - "--atomic-energy", + "--atomic", action="store_true", default=False, - help="Test the accuracy of atomic energy", + help="Test the accuracy of atomic label, i.e. energy / dipole / polar", ) # * compress model ***************************************************************** diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index 5c91547866..62687866a8 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -26,7 +26,7 @@ def test( rand_seed: Optional[int], shuffle_test: bool, detail_file: str, - atomic_energy: bool, + atomic: bool, **kwargs, ): """Test model predictions. @@ -47,7 +47,7 @@ def test( whether to shuffle tests detail_file : Optional[str] file where test details will be output - atomic_energy : bool + atomic : bool whether per atom quantities should be computed Raises @@ -83,11 +83,11 @@ def test( system, numb_test, detail_file, - atomic_energy, + atomic, append_detail=(cc != 0), ) elif dp.model_type == "dipole": - err, siz = test_dipole(dp, data, numb_test, detail_file) + err, siz = test_dipole(dp, data, numb_test, detail_file, atomic) elif dp.model_type == "polar": err, siz = test_polar(dp, data, numb_test, detail_file, global_polar=False) elif dp.model_type == "global_polar": @@ -550,6 +550,7 @@ def test_dipole( data: DeepmdData, numb_test: int, detail_file: Optional[str], + has_atom_dipole: bool, ) -> Tuple[List[np.ndarray], List[int]]: """Test energy type model. @@ -563,6 +564,8 @@ def test_dipole( munber of tests to do detail_file : Optional[str] file where test details will be output + has_atom_dipole : bool + whether atomic dipole is provided Returns ------- @@ -570,12 +573,22 @@ def test_dipole( arrays with results and their shapes """ data.add( - "dipole", 3, atomic=True, must=True, high_prec=False, type_sel=dp.get_sel_type() + "dipole", 3, atomic=has_atom_dipole, must=True, high_prec=False, type_sel=dp.get_sel_type() ) test_data = data.get_test() dipole, numb_test, _ = run_test(dp, test_data, numb_test) + + # do summation in atom dimension + if has_atom_dipole == False: + dipole = np.reshape(dipole,(dipole.shape[0], -1, 3)) + atoms = dipole.shape[1] + dipole = np.sum(dipole,axis=1) + l2f = l2err(dipole - test_data["dipole"][:numb_test]) + if has_atom_dipole == False: + l2f = l2f / atoms + log.info(f"# number of test data : {numb_test:d}") log.info(f"Dipole RMSE : {l2f:e} eV/A") diff --git a/deepmd/loss/tensor.py b/deepmd/loss/tensor.py index ceff94e14b..addccdcadf 100644 --- a/deepmd/loss/tensor.py +++ b/deepmd/loss/tensor.py @@ -12,24 +12,62 @@ class TensorLoss () : def __init__ (self, jdata, **kwarg) : try: model = kwarg['model'] - type_sel = model.get_sel_type() + self.type_sel = model.get_sel_type() except : - type_sel = None + self.type_sel = None self.tensor_name = kwarg['tensor_name'] self.tensor_size = kwarg['tensor_size'] self.label_name = kwarg['label_name'] - self.atomic = kwarg.get('atomic', True) + self.atomic = kwarg.get('atomic', None) if jdata is not None: self.scale = jdata.get('scale', 1.0) else: self.scale = 1.0 + + # YHT: added for global / local dipole combination + if self.atomic is True: # upper regulation, will control the lower behavior + self.local_weight,self.global_weight = 1.0,0.0 + elif self.atomic is False: # upper regulation, will control the lower behavior + self.local_weight,self.global_weight = 0.0,1.0 + else: # self.atomic is None, let the loss parameter decide which mode to use + if jdata is not None: + self.local_weight = jdata.get('pref_atomic_' + self.tensor_name,None) + self.global_weight = jdata.get('pref_' + self.tensor_name,None) + + # get the input parameter first + if self.local_weight is None and self.global_weight is None: + # default: downward compatibility, using local mode + self.local_weight , self.global_weight = 1.0, 0.0 + self.atomic = True + elif self.local_weight is None and self.global_weight is not None: + # using global mode only, normalize to 1 + assert self.global_weight > 0.0, AssertionError('assign a zero weight to global dipole without setting a local weight') + self.local_weight = 0.0 + self.atomic = False + elif self.local_weight is not None and self.global_weight is None: + assert self.local_weight > 0.0, AssertionError('assign a zero weight to local dipole without setting a global weight') + self.global_weight = 0.0 + self.atomic = True + else: # Both are not None + self.atomic = True if self.local_weight != 0.0 else False + assert (self.local_weight >0.0) or (self.global_weight>0.0), AssertionError('can not assian zero weight to both local and global mode') + + # normalize, not do according to Han Wang's suggestion + #temp_sum = self.local_weight + self.global_weight + #self.local_weight /= temp_sum + #self.global_weight /= temp_sum + + else: # Nothing been set, use default setting + self.local_weight,self.global_weight = 1.0,0.0 + self.atomic = True + # data required add_data_requirement(self.label_name, self.tensor_size, atomic=self.atomic, must=True, high_prec=False, - type_sel = type_sel) + type_sel = self.type_sel) def build (self, learning_rate, @@ -39,22 +77,67 @@ def build (self, 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) - more_loss = {'nonorm': l2_loss} - if not self.atomic : - atom_norm = 1./ global_cvt_2_tf_float(natoms[0]) - l2_loss = l2_loss * atom_norm + + # 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 = tf.reduce_mean( tf.square(self.scale*(polar - 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 = tf.reduce_mean( tf.square(self.scale*(global_polar - global_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)) + + # YHT: 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_more = more_loss['nonorm'] self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss)) return l2_loss, more_loss - @staticmethod - def print_header(): + def print_header(self): prop_fmt = ' %11s %11s' print_str = '' print_str += prop_fmt % ('rmse_tst', 'rmse_trn') + if self.local_weight > 0.0: + print_str += prop_fmt % ('rmse_lc_tst', 'rmse_lc_trn') + if self.global_weight > 0.0: + print_str += prop_fmt % ('rmse_gl_tst', 'rmse_gl_trn') return print_str def print_on_training(self, @@ -65,7 +148,20 @@ def print_on_training(self, feed_dict_test, feed_dict_batch) : - run_data = [self.l2_l] + # YHT: added to calculate the atoms number + atoms = 0 + if self.type_sel is not None: + for w in self.type_sel: + atoms += natoms[2+w] + else: + atoms = natoms[0] + + run_data = [self.l2_l, self.l2_more['local_loss'], self.l2_more['global_loss']] + summary_list = [self.l2_loss_summary] + if self.local_weight > 0.0: + summary_list.append(self.l2_loss_local_summary) + if self.global_weight > 0.0: + summary_list.append(self.l2_loss_global_summary) # first train data error_train = sess.run(run_data, feed_dict=feed_dict_batch) @@ -73,7 +169,8 @@ def print_on_training(self, # than test data, if tensorboard log writter is present, commpute summary # and write tensorboard logs if tb_writer: - summary_merged_op = tf.summary.merge([self.l2_loss_summary]) + #summary_merged_op = tf.summary.merge([self.l2_loss_summary]) + summary_merged_op = tf.summary.merge(summary_list) run_data.insert(0, summary_merged_op) test_out = sess.run(run_data, feed_dict=feed_dict_test) @@ -82,10 +179,14 @@ def print_on_training(self, summary = test_out.pop(0) tb_writer.add_summary(summary, cur_batch) - error_test = test_out[0] + error_test = test_out print_str = "" prop_fmt = " %11.2e %11.2e" - print_str += prop_fmt % (np.sqrt(error_test), np.sqrt(error_train)) + print_str += prop_fmt % (np.sqrt(error_test[0]), np.sqrt(error_train[0])) + if self.local_weight > 0.0: + print_str += prop_fmt % (np.sqrt(error_test[1]), np.sqrt(error_train[1]) ) + if self.global_weight > 0.0: + print_str += prop_fmt % (np.sqrt(error_test[2])/atoms, np.sqrt(error_train[2])/atoms) return print_str diff --git a/source/tests/test_argument_parser.py b/source/tests/test_argument_parser.py index 356dd430c1..a7e42ad81e 100644 --- a/source/tests/test_argument_parser.py +++ b/source/tests/test_argument_parser.py @@ -262,7 +262,7 @@ def test_parser_test(self): "--numb-test": dict(type=int, value=1), "--rand-seed": dict(type=(int, type(None)), value=12321), "--detail-file": dict(type=(str, type(None)), value="TARGET.FILE"), - "--atomic-energy": dict(type=bool), + "--atomic": dict(type=bool), } self.run_test(command="test", mapping=ARGS)