From 157f70c6a3dd2551a3c929ceeb37d2125fc3e645 Mon Sep 17 00:00:00 2001 From: root <2000011006@stu.pku.edu.cn> Date: Fri, 13 Dec 2024 18:46:49 +0800 Subject: [PATCH] Init branch --- deepmd/pt/loss/property.py | 21 +++++--- .../atomic_model/property_atomic_model.py | 16 ++---- deepmd/pt/utils/stat.py | 51 +++++++++++++------ deepmd/utils/out_stat.py | 49 ++++++++++++++++++ 4 files changed, 102 insertions(+), 35 deletions(-) diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index b59395d7f6..343208a6ac 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -85,35 +85,40 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False for property_name in self.property_name: assert label[property_name].shape == (nbz, self.property_name_dim_mapping[property_name]) concat_property.append(label[property_name]) - label["property"] = torch.cat([label["dipole_moment"],label["homo"]],dim=1) + label["property"] = torch.cat(concat_property,dim=1) assert label["property"].shape == (nbz, self.task_dim) + out_std = model.atomic_model.out_std[0][0] + out_bias = model.atomic_model.out_bias[0][0] + assert len(out_std.shape) == 1 + assert out_std.shape[0] == self.task_dim + loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0] more_loss = {} # loss if self.loss_func == "smooth_mae": loss += F.smooth_l1_loss( - label["property"], - model_pred["property"], + (label["property"]-out_bias)/out_std, + (model_pred["property"]-out_bias)/out_std, reduction="sum", beta=self.beta, ) elif self.loss_func == "mae": loss += F.l1_loss( - label["property"], model_pred["property"], reduction="sum" + (label["property"]-out_bias)/out_std, (model_pred["property"]-out_bias)/out_std, reduction="sum" ) elif self.loss_func == "mse": loss += F.mse_loss( - label["property"], - model_pred["property"], + (label["property"]-out_bias)/out_std, + (model_pred["property"]-out_bias)/out_std, reduction="sum", ) elif self.loss_func == "rmse": loss += torch.sqrt( F.mse_loss( - label["property"], - model_pred["property"], + (label["property"]-out_bias)/out_std, + (model_pred["property"]-out_bias)/out_std, reduction="mean", ) ) diff --git a/deepmd/pt/model/atomic_model/property_atomic_model.py b/deepmd/pt/model/atomic_model/property_atomic_model.py index 1fdc72b2b6..b72210a4a5 100644 --- a/deepmd/pt/model/atomic_model/property_atomic_model.py +++ b/deepmd/pt/model/atomic_model/property_atomic_model.py @@ -35,15 +35,7 @@ def apply_out_stat( The atom types. nf x nloc """ - if self.fitting_net.get_bias_method() == "normal": - out_bias, out_std = self._fetch_out_stat(self.bias_keys) - for kk in self.bias_keys: - # nf x nloc x odims, out_bias: ntypes x odims - ret[kk] = ret[kk] + out_bias[kk][atype] - return ret - elif self.fitting_net.get_bias_method() == "no_bias": - return ret - else: - raise NotImplementedError( - "Only 'normal' and 'no_bias' is supported for parameter 'bias_method'." - ) + out_bias, out_std = self._fetch_out_stat(self.bias_keys) + for kk in self.bias_keys: + ret[kk] = ret[kk] * out_std[kk][0] + out_bias[kk][0] + return ret \ No newline at end of file diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 563c006cbe..61b542c08e 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -29,6 +29,7 @@ from deepmd.utils.out_stat import ( compute_stats_from_atomic, compute_stats_from_redu, + compute_stats_property, ) from deepmd.utils.path import ( DPPath, @@ -290,6 +291,11 @@ def compute_output_stats( # remove the keys that are not in the sample keys = [keys] if isinstance(keys, str) else keys assert isinstance(keys, list) + sub_keys = [] + for key in keys: + if atomic_output.var_defs[key].sub_var_name is not None: + sub_keys.extend(atomic_output.var_defs[key].sub_var_name) + keys.extend(sub_keys) new_keys = [ ii for ii in keys @@ -297,8 +303,6 @@ def compute_output_stats( ] del keys keys = new_keys - from IPython import embed - embed() # split system based on label atomic_sampled_idx = defaultdict(list) global_sampled_idx = defaultdict(list) @@ -375,6 +379,7 @@ def compute_output_stats( # merge global/atomic bias bias_atom_e, std_atom_e = {}, {} + keys = ["property"] if ("property" in atomic_output.var_defs and (ii in keys for ii in atomic_output.var_defs["property"].sub_var_name)) else keys for kk in keys: # use atomic bias whenever available if kk in bias_atom_a: @@ -478,26 +483,42 @@ def compute_output_stats_global( std_atom_e = {} for kk in keys: if kk in stats_input: - if atomic_output is not None and atomic_output.get_data()[kk].intensive: - task_dim = stats_input[kk].shape[1] - assert merged_natoms[kk].shape == (nf[kk], ntypes) - stats_input[kk] = ( - merged_natoms[kk].sum(axis=1).reshape(-1, 1) * stats_input[kk] + if "property" in atomic_output.var_defs: + bias_atom_e[kk], std_atom_e[kk] = compute_stats_property( + stats_input[kk], + merged_natoms[kk], + assigned_bias=assigned_atom_ener[kk] + ) + else: + bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu( + stats_input[kk], + merged_natoms[kk], + assigned_bias=assigned_atom_ener[kk], + rcond=rcond, ) - assert stats_input[kk].shape == (nf[kk], task_dim) - bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu( - stats_input[kk], - merged_natoms[kk], - assigned_bias=assigned_atom_ener[kk], - rcond=rcond, - ) else: # this key does not have global labels, skip it. continue + if "property" in atomic_output.var_defs: + concat_bias = [] + concat_std = [] + for ii in atomic_output.var_defs["property"].sub_var_name: + assert ii in bias_atom_e.keys() + assert ii in std_atom_e.keys() + concat_bias.append(bias_atom_e[ii]) + concat_std.append(std_atom_e[ii]) + del bias_atom_e, std_atom_e + bias_atom_e = {} + std_atom_e = {} + bias_atom_e["property"] = np.concatenate(concat_bias, axis=-1) + std_atom_e["property"] = np.concatenate(concat_std, axis=-1) + std_atom_e["property"] = np.tile(std_atom_e["property"], (bias_atom_e["property"].shape[0], 1)) + + return bias_atom_e, std_atom_e + bias_atom_e, std_atom_e = _post_process_stat(bias_atom_e, std_atom_e) # unbias_e is only used for print rmse - if model_pred is None: unbias_e = { kk: merged_natoms[kk] @ bias_atom_e[kk].reshape(ntypes, -1) diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py index 4d0d788f8b..9ca404e063 100644 --- a/deepmd/utils/out_stat.py +++ b/deepmd/utils/out_stat.py @@ -130,3 +130,52 @@ def compute_stats_from_atomic( output[mask].std(axis=0) if output[mask].size > 0 else np.nan ) return output_bias, output_std + +def compute_stats_property( + output_redu: np.ndarray, + natoms: np.ndarray, + assigned_bias: Optional[np.ndarray] = None +) -> tuple[np.ndarray, np.ndarray]: + """Compute the output statistics. + + Given the reduced output value and the number of atoms for each atom, + compute the least-squares solution as the atomic output bias and std. + + Parameters + ---------- + output_redu + The reduced output value, shape is [nframes, *(odim0, odim1, ...)]. + natoms + The number of atoms for each atom, shape is [nframes, ntypes]. + assigned_bias + The assigned output bias, shape is [ntypes, *(odim0, odim1, ...)]. + Set to a tensor of shape (odim0, odim1, ...) filled with nan if the bias + of the type is not assigned. + rcond + Cut-off ratio for small singular values of a. + + Returns + ------- + np.ndarray + The computed output bias, shape is [ntypes, *(odim0, odim1, ...)]. + np.ndarray + The computed output std, shape is [*(odim0, odim1, ...)]. + """ + + natoms = np.array(natoms) # [nf, ntypes] + nf, ntypes = natoms.shape + output_redu = np.array(output_redu) + var_shape = list(output_redu.shape[1:]) + output_redu = output_redu.reshape(nf, -1) + # check shape + assert output_redu.ndim == 2 + assert natoms.ndim == 2 + assert output_redu.shape[0] == natoms.shape[0] # [nf,1] + + computed_output_bias = np.repeat(np.mean(output_redu,axis=0)[np.newaxis, :], ntypes, axis=0) + output_std = np.std(output_redu,axis=0) + + computed_output_bias = computed_output_bias.reshape([natoms.shape[1]] + var_shape) # noqa: RUF005 + output_std = output_std.reshape(var_shape) + + return computed_output_bias, output_std \ No newline at end of file