From 39d027e55f3c77456ae54450b956eade4d6c3249 Mon Sep 17 00:00:00 2001 From: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> Date: Sun, 7 Apr 2024 09:45:33 +0800 Subject: [PATCH] compute output stat for atomic model (#3642) This PR: - breaking change: the base atomic model is now a module. - reason: the out stat is a data attribute of the base atomic model. - implement the `compute_or_load_output_stat` for the base atomic model. the method computes both bias and std. - the derived atomic models call the `compute_or_load_output_stat` method for computing output stat. - atomic model provides the `apply_out_stat`, the derived class may override the method to define how the statistics is applied to an atomic model's output. @anyangml may need. - `out_stat` support statistics of output tensor of any shape. @iProzd please check if i took it correctly in [ce7ec1f](https://github.com/deepmodeling/deepmd-kit/pull/3642/commits/ce7ec1f39be669fcc63a05e9474bc248d9a54d8a) To be done: - atomic statistics of the bias and std. @anyangml - erialization and deserialization. --------- Signed-off-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> Co-authored-by: Han Wang Co-authored-by: Duo <50307526+iProzd@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../dpmodel/atomic_model/base_atomic_model.py | 6 + .../dpmodel/atomic_model/dp_atomic_model.py | 6 +- .../atomic_model/linear_atomic_model.py | 2 +- .../atomic_model/make_base_atomic_model.py | 3 - .../atomic_model/pairtab_atomic_model.py | 6 +- deepmd/dpmodel/output_def.py | 4 + .../model/atomic_model/base_atomic_model.py | 212 ++++++++- .../pt/model/atomic_model/dp_atomic_model.py | 16 +- .../model/atomic_model/linear_atomic_model.py | 6 +- .../atomic_model/pairtab_atomic_model.py | 21 +- deepmd/pt/model/task/invar_fitting.py | 16 +- deepmd/pt/utils/stat.py | 142 ++++-- deepmd/utils/out_stat.py | 18 +- .../common/dpmodel/test_dp_atomic_model.py | 2 +- .../tests/pt/model/test_atomic_model_stat.py | 443 ++++++++++++++++++ source/tests/pt/test_finetune.py | 15 +- source/tests/pt/test_multitask.py | 10 +- source/tests/pt/test_stat.py | 16 +- 18 files changed, 823 insertions(+), 121 deletions(-) create mode 100644 source/tests/pt/model/test_atomic_model_stat.py diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 42d1e67138..dbb344d5ca 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -27,13 +27,19 @@ class BaseAtomicModel(BaseAtomicModel_): def __init__( self, + type_map: List[str], atom_exclude_types: List[int] = [], pair_exclude_types: List[Tuple[int, int]] = [], ): super().__init__() + self.type_map = type_map self.reinit_atom_exclude(atom_exclude_types) self.reinit_pair_exclude(pair_exclude_types) + def get_type_map(self) -> List[str]: + """Get the type map.""" + return self.type_map + def reinit_atom_exclude( self, exclude_types: List[int] = [], diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 8a40f8d238..cca46d3710 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -53,7 +53,7 @@ def __init__( self.descriptor = descriptor self.fitting = fitting self.type_map = type_map - super().__init__(**kwargs) + super().__init__(type_map, **kwargs) def fitting_output_def(self) -> FittingOutputDef: """Get the output def of the fitting net.""" @@ -67,10 +67,6 @@ def get_sel(self) -> List[int]: """Get the neighbor selection.""" return self.descriptor.get_sel() - def get_type_map(self) -> List[str]: - """Get the type map.""" - return self.type_map - def mixed_types(self) -> bool: """If true, the model 1. assumes total number of atoms aligned across frames; diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 93a885f3ab..71e4aa542a 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -66,7 +66,7 @@ def __init__( self.mapping_list.append(self.remap_atype(tpmp, self.type_map)) assert len(err_msg) == 0, "\n".join(err_msg) self.mixed_types_list = [model.mixed_types() for model in self.models] - super().__init__(**kwargs) + super().__init__(type_map, **kwargs) def mixed_types(self) -> bool: """If true, the model diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py index fe5882ef36..3e02a5d076 100644 --- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py @@ -51,9 +51,6 @@ def atomic_output_def(self) -> FittingOutputDef: """ return self.fitting_output_def() - def get_output_keys(self) -> List[str]: - return list(self.atomic_output_def().keys()) - @abstractmethod def get_rcut(self) -> float: """Get the cut-off radius.""" diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index 30ab58928b..c970278bcf 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -59,9 +59,11 @@ def __init__( rcut: float, sel: Union[int, List[int]], type_map: List[str], + rcond: Optional[float] = None, + atom_ener: Optional[List[float]] = None, **kwargs, ): - super().__init__() + super().__init__(type_map, **kwargs) self.tab_file = tab_file self.rcut = rcut self.type_map = type_map @@ -69,6 +71,8 @@ def __init__( self.tab = PairTab(self.tab_file, rcut=rcut) self.type_map = type_map self.ntypes = len(type_map) + self.rcond = rcond + self.atom_ener = atom_ener if self.tab_file is not None: self.tab_info, self.tab_data = self.tab.get() diff --git a/deepmd/dpmodel/output_def.py b/deepmd/dpmodel/output_def.py index cbebb4908a..1c17fae432 100644 --- a/deepmd/dpmodel/output_def.py +++ b/deepmd/dpmodel/output_def.py @@ -224,6 +224,10 @@ def __init__( if not self.r_differentiable: raise ValueError("only r_differentiable variable can calculate hessian") + @property + def size(self): + return self.output_size + class FittingOutputDef: """Defines the shapes and other properties of the fitting network outputs. diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index 2980f0b21b..57ca21a826 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -23,6 +23,7 @@ from deepmd.pt.utils import ( AtomExcludeMask, PairExcludeMask, + env, ) from deepmd.pt.utils.nlist import ( extend_input_and_build_neighbor_list, @@ -35,19 +36,88 @@ ) log = logging.getLogger(__name__) +dtype = env.GLOBAL_PT_FLOAT_PRECISION +device = env.DEVICE BaseAtomicModel_ = make_base_atomic_model(torch.Tensor) -class BaseAtomicModel(BaseAtomicModel_): +class BaseAtomicModel(torch.nn.Module, BaseAtomicModel_): + """The base of atomic model. + + Parameters + ---------- + type_map + Mapping atom type to the name (str) of the type. + For example `type_map[1]` gives the name of the type 1. + atom_exclude_types + Exclude the atomic contribution of the given types + pair_exclude_types + Exclude the pair of atoms of the given types from computing the output + of the atomic model. Implemented by removing the pairs from the nlist. + rcond : float, optional + The condition number for the regression of atomic energy. + preset_out_bias : Dict[str, List[Optional[torch.Tensor]]], optional + Specifying atomic energy contribution in vacuum. Given by key:value pairs. + The value is a list specifying the bias. the elements can be None or np.array of output shape. + For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] + The `set_davg_zero` key in the descrptor should be set. + + """ + def __init__( self, + type_map: List[str], atom_exclude_types: List[int] = [], pair_exclude_types: List[Tuple[int, int]] = [], + rcond: Optional[float] = None, + preset_out_bias: Optional[Dict[str, torch.Tensor]] = None, ): - super().__init__() + torch.nn.Module.__init__(self) + BaseAtomicModel_.__init__(self) + self.type_map = type_map self.reinit_atom_exclude(atom_exclude_types) self.reinit_pair_exclude(pair_exclude_types) + self.rcond = rcond + self.preset_out_bias = preset_out_bias + + def init_out_stat(self): + """Initialize the output bias.""" + ntypes = self.get_ntypes() + self.bias_keys: List[str] = list(self.fitting_output_def().keys()) + self.max_out_size = max( + [self.atomic_output_def()[kk].size for kk in self.bias_keys] + ) + self.n_out = len(self.bias_keys) + out_bias_data = torch.zeros( + [self.n_out, ntypes, self.max_out_size], dtype=dtype, device=device + ) + out_std_data = torch.ones( + [self.n_out, ntypes, self.max_out_size], dtype=dtype, device=device + ) + self.register_buffer("out_bias", out_bias_data) + self.register_buffer("out_std", out_std_data) + + def __setitem__(self, key, value): + if key in ["out_bias"]: + self.out_bias = value + elif key in ["out_std"]: + self.out_std = value + else: + raise KeyError(key) + + def __getitem__(self, key): + if key in ["out_bias"]: + return self.out_bias + elif key in ["out_std"]: + return self.out_std + else: + raise KeyError(key) + + @torch.jit.export + def get_type_map(self) -> List[str]: + """Get the type map.""" + return self.type_map def reinit_atom_exclude( self, @@ -165,6 +235,7 @@ def forward_common_atomic( fparam=fparam, aparam=aparam, ) + ret_dict = self.apply_out_stat(ret_dict, atype) # nf x nloc atom_mask = ext_atom_mask[:, :nloc].to(torch.int32) @@ -210,9 +281,60 @@ def compute_or_load_stat( """ raise NotImplementedError + def compute_or_load_out_stat( + self, + merged: Union[Callable[[], List[dict]], List[dict]], + stat_file_path: Optional[DPPath] = None, + ): + """ + Compute the output statistics (e.g. energy bias) for the fitting net from packed data. + + Parameters + ---------- + merged : Union[Callable[[], List[dict]], List[dict]] + - List[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], List[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + stat_file_path : Optional[DPPath] + The path to the stat file. + + """ + self.change_out_bias( + merged, + stat_file_path=stat_file_path, + bias_adjust_mode="set-by-statistic", + ) + + def apply_out_stat( + self, + ret: Dict[str, torch.Tensor], + atype: torch.Tensor, + ): + """Apply the stat to each atomic output. + The developer may override the method to define how the bias is applied + to the atomic output of the model. + + Parameters + ---------- + ret + The returned dict by the forward_atomic method + atype + The atom types. nf x nloc + + """ + 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 + def change_out_bias( self, sample_merged, + stat_file_path: Optional[DPPath] = None, bias_adjust_mode="change-by-statistic", ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -231,22 +353,32 @@ def change_out_bias( 'change-by-statistic' : perform predictions on labels of target dataset, and do least square on the errors to obtain the target shift as bias. 'set-by-statistic' : directly use the statistic output bias in the target dataset. + stat_file_path : Optional[DPPath] + The path to the stat file. """ if bias_adjust_mode == "change-by-statistic": - delta_bias = compute_output_stats( + delta_bias, out_std = compute_output_stats( sample_merged, self.get_ntypes(), - keys=self.get_output_keys(), + keys=list(self.atomic_output_def().keys()), + stat_file_path=stat_file_path, model_forward=self._get_forward_wrapper_func(), - )["energy"] - self.set_out_bias(delta_bias, add=True) + rcond=self.rcond, + preset_bias=self.preset_out_bias, + ) + # self.set_out_bias(delta_bias, add=True) + self._store_out_stat(delta_bias, out_std, add=True) elif bias_adjust_mode == "set-by-statistic": - bias_atom = compute_output_stats( + bias_out, std_out = compute_output_stats( sample_merged, self.get_ntypes(), - keys=self.get_output_keys(), - )["energy"] - self.set_out_bias(bias_atom) + keys=list(self.atomic_output_def().keys()), + stat_file_path=stat_file_path, + rcond=self.rcond, + preset_bias=self.preset_out_bias, + ) + # self.set_out_bias(bias_out) + self._store_out_stat(bias_out, std_out) else: raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) @@ -279,3 +411,63 @@ def model_forward(coord, atype, box, fparam=None, aparam=None): return {kk: vv.detach() for kk, vv in atomic_ret.items()} return model_forward + + def _varsize( + self, + shape: List[int], + ) -> int: + output_size = 1 + len_shape = len(shape) + for i in range(len_shape): + output_size *= shape[i] + return output_size + + def _get_bias_index( + self, + kk: str, + ) -> int: + res: List[int] = [] + for i, e in enumerate(self.bias_keys): + if e == kk: + res.append(i) + assert len(res) == 1 + return res[0] + + def _store_out_stat( + self, + out_bias: Dict[str, torch.Tensor], + out_std: Dict[str, torch.Tensor], + add: bool = False, + ): + ntypes = self.get_ntypes() + out_bias_data = torch.clone(self.out_bias) + out_std_data = torch.clone(self.out_std) + for kk in out_bias.keys(): + assert kk in out_std.keys() + idx = self._get_bias_index(kk) + size = self._varsize(self.atomic_output_def()[kk].shape) + if not add: + out_bias_data[idx, :, :size] = out_bias[kk].view(ntypes, size) + else: + out_bias_data[idx, :, :size] += out_bias[kk].view(ntypes, size) + out_std_data[idx, :, :size] = out_std[kk].view(ntypes, size) + self.out_bias.copy_(out_bias_data) + self.out_std.copy_(out_std_data) + + def _fetch_out_stat( + self, + keys: List[str], + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]: + ret_bias = {} + ret_std = {} + ntypes = self.get_ntypes() + for kk in keys: + idx = self._get_bias_index(kk) + isize = self._varsize(self.atomic_output_def()[kk].shape) + ret_bias[kk] = self.out_bias[idx, :, :isize].view( + [ntypes] + list(self.atomic_output_def()[kk].shape) # noqa: RUF005 + ) + ret_std[kk] = self.out_std[idx, :, :isize].view( + [ntypes] + list(self.atomic_output_def()[kk].shape) # noqa: RUF005 + ) + return ret_bias, ret_std diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 13b8f09a79..c9c9e6ed47 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -34,7 +34,7 @@ @BaseAtomicModel.register("standard") -class DPAtomicModel(torch.nn.Module, BaseAtomicModel): +class DPAtomicModel(BaseAtomicModel): """Model give atomic prediction of some physical property. Parameters @@ -55,7 +55,7 @@ def __init__( type_map: List[str], **kwargs, ): - torch.nn.Module.__init__(self) + super().__init__(type_map, **kwargs) ntypes = len(type_map) self.type_map = type_map self.ntypes = ntypes @@ -63,9 +63,9 @@ def __init__( self.rcut = self.descriptor.get_rcut() self.sel = self.descriptor.get_sel() self.fitting_net = fitting - # order matters ntypes and type_map should be initialized first. - BaseAtomicModel.__init__(self, **kwargs) + super().init_out_stat() + @torch.jit.export def fitting_output_def(self) -> FittingOutputDef: """Get the output def of the fitting net.""" return ( @@ -79,11 +79,6 @@ def get_rcut(self) -> float: """Get the cut-off radius.""" return self.rcut - @torch.jit.export - def get_type_map(self) -> List[str]: - """Get the type map.""" - return self.type_map - def get_sel(self) -> List[int]: """Get the neighbor selection.""" return self.sel @@ -220,8 +215,7 @@ def wrapped_sampler(): return sampled self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path) - if self.fitting_net is not None: - self.fitting_net.compute_output_stats(wrapped_sampler, stat_file_path) + self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) def set_out_bias(self, out_bias: torch.Tensor, add=False) -> None: """ diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index f599399e66..f9fc97dea4 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -39,7 +39,7 @@ ) -class LinearEnergyAtomicModel(torch.nn.Module, BaseAtomicModel): +class LinearEnergyAtomicModel(BaseAtomicModel): """Linear model make linear combinations of several existing models. Parameters @@ -57,7 +57,8 @@ def __init__( type_map: List[str], **kwargs, ): - torch.nn.Module.__init__(self) + super().__init__(type_map, **kwargs) + super().init_out_stat() self.models = torch.nn.ModuleList(models) sub_model_type_maps = [md.get_type_map() for md in models] err_msg = [] @@ -78,7 +79,6 @@ def __init__( self.get_model_rcuts(), dtype=torch.float64, device=env.DEVICE ) self.nsels = torch.tensor(self.get_model_nsels(), device=env.DEVICE) - BaseAtomicModel.__init__(self, **kwargs) def mixed_types(self) -> bool: """If true, the model diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index 4db77790e9..627dffd620 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -17,9 +17,6 @@ from deepmd.pt.utils import ( env, ) -from deepmd.pt.utils.stat import ( - compute_output_stats, -) from deepmd.utils.pair_tab import ( PairTab, ) @@ -36,7 +33,7 @@ @BaseAtomicModel.register("pairtab") -class PairTabAtomicModel(torch.nn.Module, BaseAtomicModel): +class PairTabAtomicModel(BaseAtomicModel): """Pairwise tabulation energy model. This model can be used to tabulate the pairwise energy between atoms for either @@ -78,12 +75,12 @@ def __init__( atom_ener: Optional[List[float]] = None, **kwargs, ): - torch.nn.Module.__init__(self) + super().__init__(type_map, **kwargs) + super().init_out_stat() self.tab_file = tab_file self.rcut = rcut self.tab = self._set_pairtab(tab_file, rcut) - BaseAtomicModel.__init__(self, **kwargs) self.rcond = rcond self.atom_ener = atom_ener self.type_map = type_map @@ -227,17 +224,7 @@ def compute_or_load_stat( The path to the stat file. """ - bias_atom_e = compute_output_stats( - merged, - self.ntypes, - keys=["energy"], - stat_file_path=stat_file_path, - rcond=self.rcond, - atom_ener=self.atom_ener, - )["energy"] - self.bias_atom_e.copy_( - torch.tensor(bias_atom_e, device=env.DEVICE).view([self.ntypes, 1]) - ) + self.compute_or_load_out_stat(merged, stat_file_path) def set_out_bias(self, out_bias: torch.Tensor, add=False) -> None: """ diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py index 810cf2d675..01e6a8b95d 100644 --- a/deepmd/pt/model/task/invar_fitting.py +++ b/deepmd/pt/model/task/invar_fitting.py @@ -78,8 +78,11 @@ class InvarFitting(GeneralFitting): Random seed. exclude_types: List[int] Atomic contributions of the excluded atom types are set zero. - atom_ener: List[float], optional - Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. + atom_ener: List[Optional[torch.Tensor]], optional + Specifying atomic energy contribution in vacuum. + The value is a list specifying the bias. the elements can be None or np.array of output shape. + For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] + The `set_davg_zero` key in the descrptor should be set. """ @@ -100,7 +103,7 @@ def __init__( rcond: Optional[float] = None, seed: Optional[int] = None, exclude_types: List[int] = [], - atom_ener: Optional[List[float]] = None, + atom_ener: Optional[List[Optional[torch.Tensor]]] = None, **kwargs, ): self.dim_out = dim_out @@ -164,14 +167,17 @@ def compute_output_stats( The path to the stat file. """ + # [0] to get the mean (bias) bias_atom_e = compute_output_stats( merged, self.ntypes, keys=[self.var_name], stat_file_path=stat_file_path, rcond=self.rcond, - atom_ener=self.atom_ener, - )[self.var_name] + preset_bias={self.var_name: self.atom_ener} + if self.atom_ener is not None + else None, + )[0][self.var_name] self.bias_atom_e.copy_(bias_atom_e.view([self.ntypes, self.dim_out])) def output_def(self) -> FittingOutputDef: diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index d0a614cf79..d85741b231 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -2,6 +2,7 @@ import logging from typing import ( Callable, + Dict, List, Optional, Union, @@ -83,28 +84,58 @@ def _restore_from_file( keys: List[str] = ["energy"], ) -> Optional[dict]: if stat_file_path is None: - return None + return None, None stat_files = [stat_file_path / f"bias_atom_{kk}" for kk in keys] - if any(not (ii.is_file()) for ii in stat_files): - return None - ret = {} + if all(not (ii.is_file()) for ii in stat_files): + return None, None + stat_files = [stat_file_path / f"std_atom_{kk}" for kk in keys] + if all(not (ii.is_file()) for ii in stat_files): + return None, None + ret_bias = {} + ret_std = {} for kk in keys: fp = stat_file_path / f"bias_atom_{kk}" - assert fp.is_file() - ret[kk] = fp.load_numpy() - return ret + # only read the key that exists + if fp.is_file(): + ret_bias[kk] = fp.load_numpy() + for kk in keys: + fp = stat_file_path / f"std_atom_{kk}" + # only read the key that exists + if fp.is_file(): + ret_std[kk] = fp.load_numpy() + return ret_bias, ret_std def _save_to_file( stat_file_path: DPPath, - results: dict, + bias_out: dict, + std_out: dict, ): assert stat_file_path is not None stat_file_path.mkdir(exist_ok=True, parents=True) - for kk, vv in results.items(): + for kk, vv in bias_out.items(): fp = stat_file_path / f"bias_atom_{kk}" fp.save_numpy(vv) + for kk, vv in std_out.items(): + fp = stat_file_path / f"std_atom_{kk}" + fp.save_numpy(vv) + + +def _post_process_stat( + out_bias, + out_std, +): + """Post process the statistics. + + For global statistics, we do not have the std for each type of atoms, + thus fake the output std by ones for all the types. + + """ + new_std = {} + for kk, vv in out_bias.items(): + new_std[kk] = np.ones_like(vv) + return out_bias, new_std def _compute_model_predict( @@ -147,13 +178,38 @@ def model_forward_auto_batch_size(*args, **kwargs): return model_predict +def _make_preset_out_bias( + ntypes: int, + ibias: List[Optional[np.array]], +) -> Optional[np.array]: + """Make preset out bias. + + output: + a np array of shape [ntypes, *(odim0, odim1, ...)] is any item is not None + None if all items are None. + """ + if len(ibias) != ntypes: + raise ValueError("the length of preset bias list should be ntypes") + if all(ii is None for ii in ibias): + return None + for refb in ibias: + if refb is not None: + break + refb = np.array(refb) + nbias = [ + np.full_like(refb, np.nan, dtype=np.float64) if ii is None else ii + for ii in ibias + ] + return np.array(nbias) + + def compute_output_stats( merged: Union[Callable[[], List[dict]], List[dict]], ntypes: int, keys: Union[str, List[str]] = ["energy"], stat_file_path: Optional[DPPath] = None, rcond: Optional[float] = None, - atom_ener: Optional[List[float]] = None, + preset_bias: Optional[Dict[str, List[Optional[torch.Tensor]]]] = None, model_forward: Optional[Callable[..., torch.Tensor]] = None, ): """ @@ -174,8 +230,11 @@ def compute_output_stats( The path to the stat file. rcond : float, optional The condition number for the regression of atomic energy. - atom_ener : List[float], optional - Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. + preset_bias : Dict[str, List[Optional[torch.Tensor]]], optional + Specifying atomic energy contribution in vacuum. Given by key:value pairs. + The value is a list specifying the bias. the elements can be None or np.array of output shape. + For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] + The `set_davg_zero` key in the descrptor should be set. model_forward : Callable[..., torch.Tensor], optional The wrapped forward function of atomic model. If not None, the model will be utilized to generate the original energy prediction, @@ -183,7 +242,7 @@ def compute_output_stats( The difference will then be used to calculate the delta complement energy bias for each type. """ # try to restore the bias from stat file - bias_atom_e = _restore_from_file(stat_file_path, keys) + bias_atom_e, std_atom_e = _restore_from_file(stat_file_path, keys) # failed to restore the bias from stat file. compute if bias_atom_e is None: @@ -210,12 +269,16 @@ def compute_output_stats( merged_output = {kk: to_numpy_array(torch.cat(outputs[kk])) for kk in keys} # shape: (nframes, ntypes) merged_natoms = to_numpy_array(torch.cat(input_natoms)[:, 2:]) - if atom_ener is not None and len(atom_ener) > 0: - assigned_atom_ener = np.array( - [ee if ee is not None else np.nan for ee in atom_ener] - ) + nf = merged_natoms.shape[0] + if preset_bias is not None: + assigned_atom_ener = { + kk: _make_preset_out_bias(ntypes, preset_bias[kk]) + if kk in preset_bias.keys() + else None + for kk in keys + } else: - assigned_atom_ener = None + assigned_atom_ener = {kk: None for kk in keys} if model_forward is None: stats_input = merged_output @@ -224,39 +287,46 @@ def compute_output_stats( model_predict = _compute_model_predict(sampled, keys, model_forward) stats_input = {kk: merged_output[kk] - model_predict[kk] for kk in keys} - # [0]: take the first otuput (mean) of compute_stats_from_redu - bias_atom_e = { - kk: compute_stats_from_redu( + bias_atom_e = {} + std_atom_e = {} + for kk in keys: + bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu( stats_input[kk], merged_natoms, - assigned_bias=assigned_atom_ener, + assigned_bias=assigned_atom_ener[kk], rcond=rcond, - )[0] - for kk in keys - } + ) + 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_forward is None: - unbias_e = {kk: merged_natoms @ bias_atom_e[kk] for kk in keys} + unbias_e = { + kk: merged_natoms @ bias_atom_e[kk].reshape(ntypes, -1) for kk in keys + } else: unbias_e = { - kk: model_predict[kk] + merged_natoms @ bias_atom_e[kk] for kk in keys + kk: model_predict[kk].reshape(nf, -1) + + merged_natoms @ bias_atom_e[kk].reshape(ntypes, -1) + for kk in keys } atom_numbs = merged_natoms.sum(-1) + + def rmse(x): + return np.sqrt(np.mean(np.square(x))) + for kk in keys: - rmse_ae = np.sqrt( - np.mean( - np.square( - (unbias_e[kk].ravel() - merged_output[kk].ravel()) / atom_numbs - ) - ) + rmse_ae = rmse( + (unbias_e[kk].reshape(nf, -1) - merged_output[kk].reshape(nf, -1)) + / atom_numbs[:, None] ) log.info( f"RMSE of {kk} per atom after linear regression is: {rmse_ae} in the unit of {kk}." ) if stat_file_path is not None: - _save_to_file(stat_file_path, bias_atom_e) + _save_to_file(stat_file_path, bias_atom_e, std_atom_e) - ret = {kk: to_torch_tensor(bias_atom_e[kk]) for kk in keys} + ret_bias = {kk: to_torch_tensor(vv) for kk, vv in bias_atom_e.items()} + ret_std = {kk: to_torch_tensor(vv) for kk, vv in std_atom_e.items()} - return ret + return ret_bias, ret_std diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py index 3956dac654..1dcbcb1280 100644 --- a/deepmd/utils/out_stat.py +++ b/deepmd/utils/out_stat.py @@ -23,24 +23,28 @@ def compute_stats_from_redu( Parameters ---------- output_redu - The reduced output value, shape is [nframes, ndim]. + 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, ndim]. Set to nan - if not assigned. + 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, ndim]. + The computed output bias, shape is [ntypes, *(odim0, odim1, ...)]. np.ndarray - The computed output std, shape is [ntypes, ndim]. + The computed output std, shape is [*(odim0, odim1, ...)]. """ - output_redu = np.array(output_redu) natoms = np.array(natoms) + nf, _ = 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 @@ -74,6 +78,8 @@ def compute_stats_from_redu( # rest_redu: nframes, ndim rest_redu = output_redu - np.einsum("ij,jk->ik", natoms, computed_output_bias) output_std = rest_redu.std(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 diff --git a/source/tests/common/dpmodel/test_dp_atomic_model.py b/source/tests/common/dpmodel/test_dp_atomic_model.py index a3cad8b406..96bf24e451 100644 --- a/source/tests/common/dpmodel/test_dp_atomic_model.py +++ b/source/tests/common/dpmodel/test_dp_atomic_model.py @@ -41,7 +41,7 @@ def test_methods(self): md0 = DPAtomicModel(ds, ft, type_map=type_map) - self.assertEqual(md0.get_output_keys(), ["energy", "mask"]) + self.assertEqual(list(md0.atomic_output_def().keys()), ["energy", "mask"]) self.assertEqual(md0.get_type_map(), ["foo", "bar"]) self.assertEqual(md0.get_ntypes(), 2) self.assertAlmostEqual(md0.get_rcut(), self.rcut) diff --git a/source/tests/pt/model/test_atomic_model_stat.py b/source/tests/pt/model/test_atomic_model_stat.py new file mode 100644 index 0000000000..e266cf215a --- /dev/null +++ b/source/tests/pt/model/test_atomic_model_stat.py @@ -0,0 +1,443 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import tempfile +import unittest +from pathlib import ( + Path, +) +from typing import ( + Optional, +) + +import h5py +import numpy as np +import torch + +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, +) +from deepmd.pt.model.atomic_model import ( + BaseAtomicModel, + DPAtomicModel, +) +from deepmd.pt.model.descriptor.dpa1 import ( + DescrptDPA1, +) +from deepmd.pt.model.task.base_fitting import ( + BaseFitting, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.utils import ( + to_numpy_array, + to_torch_tensor, +) +from deepmd.utils.path import ( + DPPath, +) + +from .test_env_mat import ( + TestCaseSingleFrameWithNlist, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION + + +class FooFitting(torch.nn.Module, BaseFitting): + def output_def(self): + return FittingOutputDef( + [ + OutputVariableDef( + "foo", + [1], + reduciable=True, + r_differentiable=True, + c_differentiable=True, + ), + OutputVariableDef( + "pix", + [1], + reduciable=True, + r_differentiable=True, + c_differentiable=True, + ), + OutputVariableDef( + "bar", + [1, 2], + reduciable=True, + r_differentiable=True, + c_differentiable=True, + ), + ] + ) + + def serialize(self) -> dict: + raise NotImplementedError + + def forward( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + gr: Optional[torch.Tensor] = None, + g2: Optional[torch.Tensor] = None, + h2: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + ): + nf, nloc, _ = descriptor.shape + ret = {} + ret["foo"] = ( + torch.Tensor( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ) + .view([nf, nloc] + self.output_def()["foo"].shape) # noqa: RUF005 + .to(env.GLOBAL_PT_FLOAT_PRECISION) + .to(env.DEVICE) + ) + ret["pix"] = ( + torch.Tensor( + [ + [3.0, 2.0, 1.0], + [6.0, 5.0, 4.0], + ] + ) + .view([nf, nloc] + self.output_def()["pix"].shape) # noqa: RUF005 + .to(env.GLOBAL_PT_FLOAT_PRECISION) + .to(env.DEVICE) + ) + ret["bar"] = ( + torch.Tensor( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ) + .view([nf, nloc] + self.output_def()["bar"].shape) # noqa: RUF005 + .to(env.GLOBAL_PT_FLOAT_PRECISION) + .to(env.DEVICE) + ) + return ret + + +class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): + def tearDown(self): + self.tempdir.cleanup() + + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + nf, nloc, nnei = self.nlist.shape + self.merged_output_stat = [ + { + "coord": to_torch_tensor(np.zeros([2, 3, 3])), + "atype": to_torch_tensor( + np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32) + ), + "atype_ext": to_torch_tensor( + np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32) + ), + "box": to_torch_tensor(np.zeros([2, 3, 3])), + "natoms": to_torch_tensor( + np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32) + ), + # bias of foo: 1, 3 + "foo": to_torch_tensor(np.array([5.0, 7.0]).reshape(2, 1)), + # no bias of pix + # bias of bar: [1, 5], [3, 2] + "bar": to_torch_tensor( + np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2) + ), + } + ] + self.tempdir = tempfile.TemporaryDirectory() + h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) + with h5py.File(h5file, "w") as f: + pass + self.stat_file_path = DPPath(h5file, "a") + + def test_output_stat(self): + nf, nloc, nnei = self.nlist.shape + ds = DescrptDPA1( + self.rcut, + self.rcut_smth, + sum(self.sel), + self.nt, + ).to(env.DEVICE) + ft = FooFitting().to(env.DEVICE) + type_map = ["foo", "bar"] + md0 = DPAtomicModel( + ds, + ft, + type_map=type_map, + ).to(env.DEVICE) + args = [ + to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist] + ] + # nf x nloc + at = self.atype_ext[:, :nloc] + + def cvt_ret(x): + return {kk: to_numpy_array(vv) for kk, vv in x.items()} + + # 1. test run without bias + # nf x na x odim + ret0 = md0.forward_common_atomic(*args) + ret0 = cvt_ret(ret0) + expected_ret0 = {} + expected_ret0["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc] + md0.fitting_output_def()["foo"].shape) # noqa: RUF005 + expected_ret0["pix"] = np.array( + [ + [3.0, 2.0, 1.0], + [6.0, 5.0, 4.0], + ] + ).reshape([nf, nloc] + md0.fitting_output_def()["pix"].shape) # noqa: RUF005 + expected_ret0["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc] + md0.fitting_output_def()["bar"].shape) # noqa: RUF005 + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) + + # 2. test bias is applied + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret1 = md0.forward_common_atomic(*args) + ret1 = cvt_ret(ret1) + # nt x odim + foo_bias = np.array([1.0, 3.0]).reshape(2, 1) + bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) + expected_ret1 = {} + expected_ret1["foo"] = ret0["foo"] + foo_bias[at] + expected_ret1["pix"] = ret0["pix"] + expected_ret1["bar"] = ret0["bar"] + bar_bias[at] + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) + + # 3. test bias load from file + def raise_error(): + raise RuntimeError + + md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) + ret2 = md0.forward_common_atomic(*args) + ret2 = cvt_ret(ret2) + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], ret2[kk]) + + # 4. test change bias + BaseAtomicModel.change_out_bias( + md0, self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + args = [ + to_torch_tensor(ii) + for ii in [ + self.coord_ext, + to_numpy_array(self.merged_output_stat[0]["atype_ext"]), + self.nlist, + ] + ] + ret3 = md0.forward_common_atomic(*args) + ret3 = cvt_ret(ret3) + ## model output on foo: [[2, 3, 6], [5, 8, 9]] given bias [1, 3] + ## foo sumed: [11, 22] compared with [5, 7], fit target is [-6, -15] + ## fit bias is [1, -8] + ## old bias + fit bias [2, -5] + ## new model output is [[3, 4, -2], [6, 0, 1]], which sumed to [5, 7] + expected_ret3 = {} + expected_ret3["foo"] = np.array([[3, 4, -2], [6, 0, 1]]).reshape(2, 3, 1) + expected_ret3["pix"] = ret0["pix"] + for kk in ["foo", "pix"]: + np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk]) + # bar is too complicated to be manually computed. + + def test_preset_bias(self): + nf, nloc, nnei = self.nlist.shape + ds = DescrptDPA1( + self.rcut, + self.rcut_smth, + sum(self.sel), + self.nt, + ).to(env.DEVICE) + ft = FooFitting().to(env.DEVICE) + type_map = ["foo", "bar"] + preset_out_bias = { + # "foo": np.array(3.0, 2.0]).reshape(2, 1), + "foo": [None, 2], + "bar": np.array([7.0, 5.0, 13.0, 11.0]).reshape(2, 1, 2), + } + md0 = DPAtomicModel( + ds, + ft, + type_map=type_map, + preset_out_bias=preset_out_bias, + ).to(env.DEVICE) + args = [ + to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist] + ] + # nf x nloc + at = self.atype_ext[:, :nloc] + + def cvt_ret(x): + return {kk: to_numpy_array(vv) for kk, vv in x.items()} + + # 1. test run without bias + # nf x na x odim + ret0 = md0.forward_common_atomic(*args) + ret0 = cvt_ret(ret0) + expected_ret0 = {} + expected_ret0["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc] + md0.fitting_output_def()["foo"].shape) # noqa: RUF005 + expected_ret0["pix"] = np.array( + [ + [3.0, 2.0, 1.0], + [6.0, 5.0, 4.0], + ] + ).reshape([nf, nloc] + md0.fitting_output_def()["pix"].shape) # noqa: RUF005 + expected_ret0["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc] + md0.fitting_output_def()["bar"].shape) # noqa: RUF005 + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) + + # 2. test bias is applied + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret1 = md0.forward_common_atomic(*args) + ret1 = cvt_ret(ret1) + # foo sums: [5, 7], + # given bias of type 1 being 2, the bias left for type 0 is [5-2*1, 7-2*2] = [3,3] + # the solution of type 0 is 1.8 + foo_bias = np.array([1.8, preset_out_bias["foo"][1]]).reshape(2, 1) + bar_bias = preset_out_bias["bar"] + expected_ret1 = {} + expected_ret1["foo"] = ret0["foo"] + foo_bias[at] + expected_ret1["pix"] = ret0["pix"] + expected_ret1["bar"] = ret0["bar"] + bar_bias[at] + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) + + # 3. test bias load from file + def raise_error(): + raise RuntimeError + + md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) + ret2 = md0.forward_common_atomic(*args) + ret2 = cvt_ret(ret2) + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], ret2[kk]) + + # 4. test change bias + BaseAtomicModel.change_out_bias( + md0, self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + args = [ + to_torch_tensor(ii) + for ii in [ + self.coord_ext, + to_numpy_array(self.merged_output_stat[0]["atype_ext"]), + self.nlist, + ] + ] + ret3 = md0.forward_common_atomic(*args) + ret3 = cvt_ret(ret3) + ## model output on foo: [[2.8, 3.8, 5], [5.8, 7., 8.]] given bias [1.8, 2] + ## foo sumed: [11.6, 20.8] compared with [5, 7], fit target is [-6.6, -13.8] + ## fit bias is [-7, 2] (2 is assigned. -7 is fit to [-8.6, -17.8]) + ## old bias[1.8,2] + fit bias[-7, 2] = [-5.2, 4] + ## new model output is [[-4.2, -3.2, 7], [-1.2, 9, 10]] + expected_ret3 = {} + expected_ret3["foo"] = np.array([[-4.2, -3.2, 7.0], [-1.2, 9.0, 10.0]]).reshape( + 2, 3, 1 + ) + expected_ret3["pix"] = ret0["pix"] + for kk in ["foo", "pix"]: + np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk]) + # bar is too complicated to be manually computed. + + def test_preset_bias_all_none(self): + nf, nloc, nnei = self.nlist.shape + ds = DescrptDPA1( + self.rcut, + self.rcut_smth, + sum(self.sel), + self.nt, + ).to(env.DEVICE) + ft = FooFitting().to(env.DEVICE) + type_map = ["foo", "bar"] + preset_out_bias = { + "foo": [None, None], + } + md0 = DPAtomicModel( + ds, + ft, + type_map=type_map, + preset_out_bias=preset_out_bias, + ).to(env.DEVICE) + args = [ + to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist] + ] + # nf x nloc + at = self.atype_ext[:, :nloc] + + def cvt_ret(x): + return {kk: to_numpy_array(vv) for kk, vv in x.items()} + + # 1. test run without bias + # nf x na x odim + ret0 = md0.forward_common_atomic(*args) + ret0 = cvt_ret(ret0) + expected_ret0 = {} + expected_ret0["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc] + md0.fitting_output_def()["foo"].shape) # noqa: RUF005 + expected_ret0["pix"] = np.array( + [ + [3.0, 2.0, 1.0], + [6.0, 5.0, 4.0], + ] + ).reshape([nf, nloc] + md0.fitting_output_def()["pix"].shape) # noqa: RUF005 + expected_ret0["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc] + md0.fitting_output_def()["bar"].shape) # noqa: RUF005 + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) + + # 2. test bias is applied + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret1 = md0.forward_common_atomic(*args) + ret1 = cvt_ret(ret1) + # nt x odim + foo_bias = np.array([1.0, 3.0]).reshape(2, 1) + bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) + expected_ret1 = {} + expected_ret1["foo"] = ret0["foo"] + foo_bias[at] + expected_ret1["pix"] = ret0["pix"] + expected_ret1["bar"] = ret0["bar"] + bar_bias[at] + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) diff --git a/source/tests/pt/test_finetune.py b/source/tests/pt/test_finetune.py index b7120414ba..8f299ce542 100644 --- a/source/tests/pt/test_finetune.py +++ b/source/tests/pt/test_finetune.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import tempfile import unittest -from copy import ( - deepcopy, -) from pathlib import ( Path, ) @@ -41,11 +38,9 @@ class FinetuneTest: def test_finetune_change_out_bias(self): # get model model = get_model(self.model_config) - fitting_net = model.get_fitting_net() - fitting_net["bias_atom_e"] = torch.rand_like(fitting_net["bias_atom_e"]) - energy_bias_before = deepcopy( - to_numpy_array(fitting_net["bias_atom_e"]).reshape(-1) - ) + atomic_model = model.atomic_model + atomic_model["out_bias"] = torch.rand_like(atomic_model["out_bias"]) + energy_bias_before = to_numpy_array(atomic_model["out_bias"])[0].ravel() # prepare original model for test dp = torch.jit.script(model) @@ -60,9 +55,7 @@ def test_finetune_change_out_bias(self): self.sampled, bias_adjust_mode="change-by-statistic", ) - energy_bias_after = deepcopy( - to_numpy_array(fitting_net["bias_atom_e"]).reshape(-1) - ) + energy_bias_after = to_numpy_array(atomic_model["out_bias"])[0].ravel() # get ground-truth energy bias change sorter = np.argsort(full_type_map) diff --git a/source/tests/pt/test_multitask.py b/source/tests/pt/test_multitask.py index 8bdb42df52..3c78484e1f 100644 --- a/source/tests/pt/test_multitask.py +++ b/source/tests/pt/test_multitask.py @@ -128,17 +128,21 @@ def test_multitask_train(self): multi_state_dict[state_key], multi_state_dict_finetuned[state_key], ) - elif "model_2" in state_key and "bias_atom_e" not in state_key: + elif "model_2" in state_key and "out_bias" not in state_key: torch.testing.assert_close( multi_state_dict[state_key], multi_state_dict_finetuned[state_key], ) - elif "model_3" in state_key and "bias_atom_e" not in state_key: + elif "model_3" in state_key and "out_bias" not in state_key: torch.testing.assert_close( multi_state_dict[state_key.replace("model_3", "model_2")], multi_state_dict_finetuned[state_key], ) - elif "model_4" in state_key and "fitting_net" not in state_key: + elif ( + "model_4" in state_key + and "fitting_net" not in state_key + and "out_bias" not in state_key + ): torch.testing.assert_close( multi_state_dict[state_key.replace("model_4", "model_2")], multi_state_dict_finetuned[state_key], diff --git a/source/tests/pt/test_stat.py b/source/tests/pt/test_stat.py index 2362821dfa..76549f0c7d 100644 --- a/source/tests/pt/test_stat.py +++ b/source/tests/pt/test_stat.py @@ -366,12 +366,12 @@ def test_calc_and_load(self): type_map = self.type_map # compute from sample - ret0 = compute_output_stats( + ret0, _ = compute_output_stats( self.sampled, len(type_map), keys=["energy"], stat_file_path=stat_file_path, - atom_ener=None, + preset_bias=None, model_forward=None, ) # ground truth @@ -394,12 +394,12 @@ def raise_error(): # hack!!! # suppose to load stat from file, if from sample, an error will raise. - ret1 = compute_output_stats( + ret1, _ = compute_output_stats( raise_error, len(type_map), keys=["energy"], stat_file_path=stat_file_path, - atom_ener=None, + preset_bias=None, model_forward=None, ) np.testing.assert_almost_equal( @@ -407,21 +407,21 @@ def raise_error(): ) def test_assigned(self): - atom_ener = np.array([3.0, 5.0]).reshape(2, 1) + atom_ener = {"energy": np.array([3.0, 5.0]).reshape(2, 1)} stat_file_path = self.stat_file_path type_map = self.type_map # from assigned atom_ener - ret2 = compute_output_stats( + ret2, _ = compute_output_stats( self.sampled, len(type_map), keys=["energy"], stat_file_path=stat_file_path, - atom_ener=atom_ener, + preset_bias=atom_ener, model_forward=None, ) np.testing.assert_almost_equal( - to_numpy_array(ret2["energy"]), atom_ener, decimal=10 + to_numpy_array(ret2["energy"]), atom_ener["energy"], decimal=10 )