diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 25b770a5fb..720833b5cd 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -51,6 +51,7 @@ check_preset_out_bias, normalize_preset_out_bias, preset_out_bias_rows, + recover_nonfinite_preset_bias, remap_preset_out_bias, ) @@ -822,6 +823,7 @@ def change_out_bias( stat_file_path: DPPath | None = None, bias_adjust_mode: str = "change-by-statistic", observed_type: list[str] | None = None, + preset_out_bias: dict | None = None, ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -845,6 +847,14 @@ def change_out_bias( ('set-by-statistic') or their stored bias ('change-by-statistic'); no statistics are computed, read or written for such an output, whose output std keeps its stored value. + preset_out_bias : dict, optional + An explicit downstream preset. Assigned outputs use this bias and + skip fitting; other outputs fit residuals against the stored bias. + The model preset is used when omitted, preserving training and + fine-tuning behavior. Pass an empty dict to fit from the data + without reapplying the model's stored preset. + Before residual fitting, non-finite stored cells covered by the + checkpoint preset are recovered; finite adjusted cells are retained. stat_file_path : Optional[DPPath] The path to the stat file. observed_type : list[str], optional @@ -859,17 +869,32 @@ def change_out_bias( if bias_adjust_mode not in ("change-by-statistic", "set-by-statistic"): raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) change = bias_adjust_mode == "change-by-statistic" + preset = ( + normalize_preset_out_bias(preset_out_bias, self.type_map) + if preset_out_bias is not None + else self.preset_out_bias + ) distinguish_types = self.get_compute_stats_distinguish_types() - check_preset_out_bias(self.preset_out_bias, self.bias_keys, distinguish_types) + check_preset_out_bias(preset, self.bias_keys, distinguish_types) sampled = sample_merged + # A per-call override must not propagate a corrupt checkpoint baseline. + if change and preset_out_bias is not None and self.preset_out_bias: + recovered = recover_nonfinite_preset_bias( + self.preset_out_bias, + self.type_map, + to_numpy_array(self.out_bias), + self.bias_keys, + [self.atomic_output_def()[kk].size for kk in self.bias_keys], + ) + self._store_out_stat(recovered) # === Step 1. Outputs fixed by the preset === fitted_keys = self.bias_keys - if self.preset_out_bias: + if preset: if observed_type is None: sampled = sample_merged() if callable(sample_merged) else sample_merged observed_type = collect_observed_types(sampled, self.type_map) rows = preset_out_bias_rows( - self.preset_out_bias, + preset, self.type_map, observed_type, to_numpy_array(self.out_bias), diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index b74fb6b697..607f7dfb83 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -926,6 +926,7 @@ def change_out_bias( self, merged: Any, bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -936,8 +937,17 @@ def change_out_bias( bias_adjust_mode : str The mode for changing output bias: 'change-by-statistic' or 'set-by-statistic'. + preset_out_bias : dict, optional + Per-call preset override. Omit to use the configured preset; + pass an empty dict to fit without fixing preset outputs. + Non-finite stored cells covered by the checkpoint preset are + recovered before residual fitting; finite values are retained. """ - self.atomic_model.change_out_bias(merged, bias_adjust_mode=bias_adjust_mode) + self.atomic_model.change_out_bias( + merged, + bias_adjust_mode=bias_adjust_mode, + preset_out_bias=preset_out_bias, + ) def _input_type_cast( self, diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 8678e51398..ae7ed8bae3 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -31,6 +31,10 @@ from deepmd.dpmodel.output_def import ( ModelOutputDef, ) +from deepmd.utils.preset_out_bias import ( + normalize_preset_out_bias, + remap_preset_out_bias, +) from deepmd.utils.spin import ( Spin, ) @@ -603,6 +607,7 @@ def change_out_bias( self, merged: Callable[[], list[dict]] | list[dict], bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> None: """Change the output bias of atomic model according to the input data and the pretrained model. @@ -620,13 +625,24 @@ 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. + preset_out_bias : dict, optional + Preset on the public real-atom type map. It is normalized and + extended with unassigned virtual types before forwarding. Omit + to retain the model preset; an empty dict fits without fixing + preset outputs. Non-finite stored cells covered by a checkpoint + preset are recovered before residual fitting. """ + preset_out_bias = remap_preset_out_bias( + normalize_preset_out_bias(preset_out_bias, self.get_type_map()), + list(range(self.get_ntypes())) + [-1] * self.get_ntypes(), + ) spin_sampled_func = self._get_spin_sampled_func( merged if callable(merged) else lambda: merged ) self.backbone_model.change_out_bias( spin_sampled_func, bias_adjust_mode=bias_adjust_mode, + preset_out_bias=preset_out_bias, ) def change_type_map( diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py index 22cb543ced..20757fae9f 100644 --- a/deepmd/dpmodel/train/trainer.py +++ b/deepmd/dpmodel/train/trainer.py @@ -74,12 +74,17 @@ def change_model_out_bias( *, bias_adjust_mode: str = "change-by-statistic", recompute_input_stats: bool = False, + preset_out_bias: dict | None = None, ) -> Any: """Change one model's output bias and log the before/after values.""" old_bias = deepcopy(model.get_out_bias()) + preset_kwargs = ( + {"preset_out_bias": preset_out_bias} if preset_out_bias is not None else {} + ) model.change_out_bias( sample_func, bias_adjust_mode=bias_adjust_mode, + **preset_kwargs, ) new_bias = deepcopy(model.get_out_bias()) diff --git a/deepmd/pd/entrypoints/main.py b/deepmd/pd/entrypoints/main.py index cf15e99192..237c860125 100644 --- a/deepmd/pd/entrypoints/main.py +++ b/deepmd/pd/entrypoints/main.py @@ -552,7 +552,10 @@ def change_bias( nbatches, ) updated_model = training.model_change_out_bias( - model_to_change, sampled_data, _bias_adjust_mode=bias_adjust_mode + model_to_change, + sampled_data, + _bias_adjust_mode=bias_adjust_mode, + preset_out_bias={}, ) if not multi_task: diff --git a/deepmd/pd/model/atomic_model/base_atomic_model.py b/deepmd/pd/model/atomic_model/base_atomic_model.py index 8a8aaefd48..d3fdaed5fe 100644 --- a/deepmd/pd/model/atomic_model/base_atomic_model.py +++ b/deepmd/pd/model/atomic_model/base_atomic_model.py @@ -48,6 +48,7 @@ check_preset_out_bias, normalize_preset_out_bias, preset_out_bias_rows, + recover_nonfinite_preset_bias, remap_preset_out_bias, ) @@ -501,6 +502,7 @@ def change_out_bias( stat_file_path: DPPath | None = None, bias_adjust_mode: str = "change-by-statistic", observed_type: list[str] | None = None, + preset_out_bias: dict | None = None, ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -524,6 +526,14 @@ def change_out_bias( ('set-by-statistic') or their stored bias ('change-by-statistic'); no statistics are computed, read or written for such an output, whose output std keeps its stored value. + preset_out_bias : dict, optional + An explicit downstream preset. Assigned outputs use this bias and + skip fitting; other outputs fit residuals against the stored bias. + The model preset is used when omitted, preserving training and + fine-tuning behavior. Pass an empty dict to fit from the data + without reapplying the model's stored preset. + Before residual fitting, non-finite stored cells covered by the + checkpoint preset are recovered; finite adjusted cells are retained. stat_file_path : Optional[DPPath] The path to the stat file. observed_type : list[str], optional @@ -533,17 +543,34 @@ def change_out_bias( if bias_adjust_mode not in ("change-by-statistic", "set-by-statistic"): raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) change = bias_adjust_mode == "change-by-statistic" + preset = ( + normalize_preset_out_bias(preset_out_bias, self.type_map) + if preset_out_bias is not None + else self.preset_out_bias + ) distinguish_types = self.get_compute_stats_distinguish_types() - check_preset_out_bias(self.preset_out_bias, self.bias_keys, distinguish_types) + check_preset_out_bias(preset, self.bias_keys, distinguish_types) sampled = sample_merged + # A per-call override must not propagate a corrupt checkpoint baseline. + if change and preset_out_bias is not None and self.preset_out_bias: + recovered = recover_nonfinite_preset_bias( + self.preset_out_bias, + self.type_map, + to_numpy_array(self.out_bias), + self.bias_keys, + [self.atomic_output_def()[kk].size for kk in self.bias_keys], + ) + self._store_out_stat( + {kk: to_paddle_tensor(vv) for kk, vv in recovered.items()} + ) # === Step 1. Outputs fixed by the preset === fitted_keys = self.bias_keys - if self.preset_out_bias: + if preset: if observed_type is None: sampled = sample_merged() if callable(sample_merged) else sample_merged observed_type = collect_observed_types(sampled, self.type_map) rows = preset_out_bias_rows( - self.preset_out_bias, + preset, self.type_map, observed_type, to_numpy_array(self.out_bias), diff --git a/deepmd/pd/model/model/make_model.py b/deepmd/pd/model/model/make_model.py index f5beb501fc..3fa293cf64 100644 --- a/deepmd/pd/model/model/make_model.py +++ b/deepmd/pd/model/model/make_model.py @@ -231,6 +231,7 @@ def change_out_bias( self, merged: list[dict] | Callable[[], list[dict]], bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> None: """Change the output bias of atomic model according to the input data and the pretrained model. @@ -248,10 +249,16 @@ 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. + preset_out_bias : dict, optional + Per-call preset override. Omit to use the configured preset; + pass an empty dict to fit without fixing preset outputs. + Non-finite stored cells covered by the checkpoint preset are + recovered before residual fitting; finite values are retained. """ self.atomic_model.change_out_bias( merged, bias_adjust_mode=bias_adjust_mode, + preset_out_bias=preset_out_bias, ) def forward_common_lower( diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py index a33908c632..7db7d89895 100644 --- a/deepmd/pd/train/training.py +++ b/deepmd/pd/train/training.py @@ -1455,11 +1455,16 @@ def model_change_out_bias( _model: Any, _sample_func: Any, _bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> None: old_bias = deepcopy(_model.get_out_bias()) + preset_kwargs = ( + {"preset_out_bias": preset_out_bias} if preset_out_bias is not None else {} + ) _model.change_out_bias( _sample_func, bias_adjust_mode=_bias_adjust_mode, + **preset_kwargs, ) new_bias = deepcopy(_model.get_out_bias()) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 9d54519e9d..d3a3b53901 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -637,7 +637,10 @@ def change_bias( min_pair_dist=0.0 if pair_filter is None else float(pair_filter.default), ) updated_model = training.model_change_out_bias( - model_to_change, sampled_data, _bias_adjust_mode=bias_adjust_mode + model_to_change, + sampled_data, + _bias_adjust_mode=bias_adjust_mode, + preset_out_bias={}, ) if not multi_task: diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index d18e780cd3..5497f33fe3 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -52,6 +52,7 @@ check_preset_out_bias, normalize_preset_out_bias, preset_out_bias_rows, + recover_nonfinite_preset_bias, remap_preset_out_bias, ) @@ -629,6 +630,7 @@ def change_out_bias( bias_adjust_mode: str = "change-by-statistic", model_forward: Callable[..., dict[str, torch.Tensor]] | None = None, observed_type: list[str] | None = None, + preset_out_bias: dict | None = None, ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -654,6 +656,14 @@ def change_out_bias( ('set-by-statistic') or their stored bias ('change-by-statistic'); no statistics are computed, read or written for such an output, whose output std keeps its stored value. + preset_out_bias : dict, optional + An explicit downstream preset. Assigned outputs use this bias and + skip fitting; other outputs fit residuals against the stored bias. + The model preset is used when omitted, preserving training and + fine-tuning behavior. Pass an empty dict to fit from the data + without reapplying the model's stored preset. + Before residual fitting, non-finite stored cells covered by the + checkpoint preset are recovered; finite adjusted cells are retained. model_forward : Callable[..., dict[str, torch.Tensor]], optional Predictor of the complete atomic outputs used by 'change-by-statistic'. Defaults to the forward of this atomic model; a model that adds @@ -665,17 +675,34 @@ def change_out_bias( if bias_adjust_mode not in ("change-by-statistic", "set-by-statistic"): raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) change = bias_adjust_mode == "change-by-statistic" + preset = ( + normalize_preset_out_bias(preset_out_bias, self.type_map) + if preset_out_bias is not None + else self.preset_out_bias + ) distinguish_types = self.get_compute_stats_distinguish_types() - check_preset_out_bias(self.preset_out_bias, self.bias_keys, distinguish_types) + check_preset_out_bias(preset, self.bias_keys, distinguish_types) sampled = sample_merged + # A per-call override must not propagate a corrupt checkpoint baseline. + if change and preset_out_bias is not None and self.preset_out_bias: + recovered = recover_nonfinite_preset_bias( + self.preset_out_bias, + self.type_map, + to_numpy_array(self.out_bias), + self.bias_keys, + [self.atomic_output_def()[kk].size for kk in self.bias_keys], + ) + self._store_out_stat( + {kk: to_torch_tensor(vv) for kk, vv in recovered.items()} + ) # === Step 1. Outputs fixed by the preset === fitted_keys = self.bias_keys - if self.preset_out_bias: + if preset: if observed_type is None: sampled = sample_merged() if callable(sample_merged) else sample_merged observed_type = collect_observed_types(sampled, self.type_map) rows = preset_out_bias_rows( - self.preset_out_bias, + preset, self.type_map, observed_type, to_numpy_array(self.out_bias), diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 185122349d..9b83dc8476 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -335,6 +335,7 @@ def change_out_bias( self, merged: Any, bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> None: """Change the output bias of atomic model according to the input data and the pretrained model. @@ -355,10 +356,16 @@ 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. + preset_out_bias : dict, optional + Per-call preset override. Omit to use the configured preset; + pass an empty dict to fit without fixing preset outputs. + Non-finite stored cells covered by the checkpoint preset are + recovered before residual fitting; finite values are retained. """ self.atomic_model.change_out_bias( merged, bias_adjust_mode=bias_adjust_mode, + preset_out_bias=preset_out_bias, model_forward=self.predict_atomic_outputs_for_stat, ) diff --git a/deepmd/pt/model/model/sezm_spin_model.py b/deepmd/pt/model/model/sezm_spin_model.py index 86b7eeca3d..245bd7a5f6 100644 --- a/deepmd/pt/model/model/sezm_spin_model.py +++ b/deepmd/pt/model/model/sezm_spin_model.py @@ -43,6 +43,10 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.preset_out_bias import ( + normalize_preset_out_bias, + remap_preset_out_bias, +) from deepmd.utils.spin import ( Spin, ) @@ -430,14 +434,35 @@ def change_out_bias( self, merged: Callable[[], list[dict[str, Any]]] | list[dict[str, Any]], bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> None: - """Change output bias using spin-expanded sampled data.""" + """Change output bias using spin-expanded sampled data. + + Parameters + ---------- + merged : callable or list[dict] + Samples containing real-atom coordinates, types and spins. + bias_adjust_mode : str + 'change-by-statistic' fits residual shifts; 'set-by-statistic' + fits absolute biases. + preset_out_bias : dict, optional + Preset on the public real-atom type map. It is normalized and + extended with unassigned virtual types before forwarding. Omit + to retain the model preset; an empty dict fits without fixing + preset outputs. Non-finite stored cells covered by a checkpoint + preset are recovered before residual fitting. + """ + preset_out_bias = remap_preset_out_bias( + normalize_preset_out_bias(preset_out_bias, self.get_type_map()), + list(range(self.get_ntypes())) + [-1] * self.get_ntypes(), + ) spin_sampled_func = self._get_spin_sampled_func( merged if callable(merged) else lambda: merged ) super().change_out_bias( spin_sampled_func, bias_adjust_mode=bias_adjust_mode, + preset_out_bias=preset_out_bias, ) def change_type_map( diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index 550de0bfa4..cb562bf7fc 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -27,6 +27,10 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.preset_out_bias import ( + normalize_preset_out_bias, + remap_preset_out_bias, +) from deepmd.utils.spin import ( Spin, ) @@ -518,6 +522,7 @@ def change_out_bias( self, merged: Callable[[], list[dict]] | list[dict], bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> None: """Change the output bias of atomic model according to the input data and the pretrained model. @@ -535,13 +540,24 @@ 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. + preset_out_bias : dict, optional + Preset on the public real-atom type map. It is normalized and + extended with unassigned virtual types before forwarding. Omit + to retain the model preset; an empty dict fits without fixing + preset outputs. Non-finite stored cells covered by a checkpoint + preset are recovered before residual fitting. """ + preset_out_bias = remap_preset_out_bias( + normalize_preset_out_bias(preset_out_bias, self.get_type_map()), + list(range(self.get_ntypes())) + [-1] * self.get_ntypes(), + ) spin_sampled_func = self._get_spin_sampled_func( merged if callable(merged) else lambda: merged ) self.backbone_model.change_out_bias( spin_sampled_func, bias_adjust_mode=bias_adjust_mode, + preset_out_bias=preset_out_bias, ) def change_type_map( diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 7efc76e4d9..a808c0dca1 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -2603,6 +2603,7 @@ def model_change_out_bias( _model: Any, _sample_func: Callable[[], Any], _bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> Any: from deepmd.pt.model.model.dp_model import ( DPModelCommon, @@ -2613,4 +2614,5 @@ def model_change_out_bias( _sample_func, bias_adjust_mode=_bias_adjust_mode, recompute_input_stats=isinstance(_model, DPModelCommon), + preset_out_bias=preset_out_bias, ) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index e4ff2f813d..83bf08cc98 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -852,7 +852,10 @@ def change_bias( 0.0 if pair_filter is None else float(pair_filter.default), ) model_to_change = model_change_out_bias( - model_to_change, sampled_data, _bias_adjust_mode=bias_adjust_mode + model_to_change, + sampled_data, + _bias_adjust_mode=bias_adjust_mode, + preset_out_bias={}, ) if input_file.endswith(".pt"): diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f1e7f27c58..22fbc3e97c 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -3571,6 +3571,7 @@ def model_change_out_bias( _model: Any, _sample_func: Any, _bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> Any: """Change the output bias of a model based on sampled data. @@ -3582,6 +3583,9 @@ def model_change_out_bias( Callable that returns sampled data for bias computation. _bias_adjust_mode ``"change-by-statistic"`` or ``"set-by-statistic"``. + preset_out_bias + Optional preset override. An empty mapping ignores checkpoint presets; + omitting it preserves training and fine-tuning behavior. Returns ------- @@ -3596,6 +3600,7 @@ def model_change_out_bias( _sample_func, bias_adjust_mode=_bias_adjust_mode, recompute_input_stats=isinstance(_model, DPModelCommon), + preset_out_bias=preset_out_bias, ) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 9f9c98c26c..8a7131e4f3 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3651,7 +3651,9 @@ def model_args( "Elements outside the `type_map` are ignored and every element that occurs in the data must be assigned. " "A table is resolved once when the input is processed and its values are stored in the model. An assigned " "output is taken from the preset without statistics, both when a model is initialized and when its bias is " - "changed by fine-tuning or `dp change-bias`. Dipole models apply no output bias and take no preset." + "changed by fine-tuning. `dp change-bias` fits from its input data without reapplying the model's stored " + "preset; explicit `--bias-value` values override the bias directly. " + "Dipole models apply no output bias and take no preset." ) doc_finetune_head = ( "The chosen fitting net to fine-tune on, when doing multi-task fine-tuning. " diff --git a/deepmd/utils/preset_out_bias.py b/deepmd/utils/preset_out_bias.py index 48a53f644a..f5bc55aaed 100644 --- a/deepmd/utils/preset_out_bias.py +++ b/deepmd/utils/preset_out_bias.py @@ -401,6 +401,48 @@ def preset_out_bias_rows( return rows +def recover_nonfinite_preset_bias( + preset_out_bias: dict[str, list[list | None]], + type_map: list[str], + stored_bias: np.ndarray, + keys: list[str], + sizes: list[int], +) -> dict[str, np.ndarray]: + """Recover non-finite stored cells from presets without fixing fitted outputs. + + Parameters + ---------- + preset_out_bias + Normalized checkpoint preset. Unassigned cells cannot be recovered. + type_map + Element names in stored type order, including types absent from data. + stored_bias + Stored bias with shape (n_out, ntypes, max_size). + keys + Output names in stored order. + sizes + Flattened size of each output. + + Returns + ------- + dict[str, np.ndarray] + Changed output rows. All finite stored cells are preserved. + """ + preset_rows = preset_out_bias_rows( + preset_out_bias, type_map, [], stored_bias, keys, sizes, keep_unassigned=True + ) + recovered = {} + for index, (key, size) in enumerate(zip(keys, sizes, strict=True)): + if key not in preset_rows: + continue + stored = stored_bias[index, :, :size] + replacement = preset_rows[key] + mask = ~np.isfinite(stored) & np.isfinite(replacement) + if np.any(mask): + recovered[key] = np.where(mask, replacement, stored) + return recovered + + def override_assigned_bias( bias: np.ndarray, assigned_bias: np.ndarray | None, diff --git a/doc/model/train-energy.md b/doc/model/train-energy.md index 83bf693866..d97d99967a 100644 --- a/doc/model/train-energy.md +++ b/doc/model/train-energy.md @@ -211,9 +211,13 @@ the autograd route instead. `vacuum_ref` is available for DPA4/SeZM on the PyTorch backend and for the graph-native models (DPA1, DPA2, DPA4 and DPA4C) on the pt_expt backend, whose -flat node axis carries the reference atoms. Fine-tuning and `dp change-bias` -keep the reference; a `preset_out_bias` given to them fixes the bias again -without statistics. +flat node axis carries the reference atoms. Fine-tuning keeps the reference +and applies its configured `preset_out_bias` without fitting assigned outputs. +Standalone `dp change-bias` also keeps the reference, but fits residuals from +the stored bias without reapplying checkpoint presets. Non-finite stored bias +cells covered by a finite checkpoint preset are recovered before residual +fitting; finite adjusted values are preserved. To set an energy bias explicitly +without statistics, use `dp change-bias --bias-value ...`. ## Loss diff --git a/source/tests/common/dpmodel/test_atomic_model_global_stat.py b/source/tests/common/dpmodel/test_atomic_model_global_stat.py index 0059374e0a..1ddc7d1acc 100644 --- a/source/tests/common/dpmodel/test_atomic_model_global_stat.py +++ b/source/tests/common/dpmodel/test_atomic_model_global_stat.py @@ -7,6 +7,9 @@ from typing import ( NoReturn, ) +from unittest.mock import ( + patch, +) import h5py import numpy as np @@ -330,6 +333,41 @@ def test_preset_bias(self) -> None: np.testing.assert_almost_equal(out_bias["foo"], foo_bias) np.testing.assert_almost_equal(out_bias["bar"], bar_bias) + # Standalone change-bias explicitly ignores the checkpoint preset. + adjusted_bias = foo_bias + 7.0 + md0._store_out_stat({"foo": adjusted_bias}) + forward = md0._get_forward_wrapper_func() + seen_biases = [] + + def check_forward(*args, **kwargs): + bias, _ = md0._fetch_out_stat(["foo"]) + seen_biases.append(bias["foo"].copy()) + return forward(*args, **kwargs) + + with patch.object(md0, "_get_forward_wrapper_func", return_value=check_forward): + md0.change_out_bias( + self.merged_output_stat, + bias_adjust_mode="change-by-statistic", + preset_out_bias={}, + ) + self.assertTrue(seen_biases) + for seen_bias in seen_biases: + np.testing.assert_almost_equal(seen_bias, adjusted_bias) + out_bias, _ = md0._fetch_out_stat(["foo", "bar"]) + np.testing.assert_almost_equal( + out_bias["foo"], np.array([2.0, -5.0]).reshape(2, 1) + ) + assert not np.allclose(out_bias["bar"], bar_bias) + + # The override does not remove the model's configured preset. + np.testing.assert_almost_equal(md0.preset_out_bias["foo"], foo_bias) + md0.change_out_bias( + self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + out_bias, _ = md0._fetch_out_stat(["foo", "bar"]) + np.testing.assert_almost_equal(out_bias["foo"], foo_bias) + np.testing.assert_almost_equal(out_bias["bar"], bar_bias) + # 4. a preset that leaves an observed type unassigned is rejected md1 = DPDPAtomicModel( ds, diff --git a/source/tests/common/dpmodel/test_spin_model_virtual_types.py b/source/tests/common/dpmodel/test_spin_model_virtual_types.py index 727eac3b53..4795dfb971 100644 --- a/source/tests/common/dpmodel/test_spin_model_virtual_types.py +++ b/source/tests/common/dpmodel/test_spin_model_virtual_types.py @@ -80,6 +80,18 @@ def test_dense_spin_expansion_preserves_virtual_types(model) -> None: ) +def test_real_type_sequence_bias_override(model) -> None: + sample = { + "coord": np.arange(9, dtype=np.float64).reshape(1, 3, 3), + "atype": np.array([[0, 1, 2]], dtype=np.int32), + "spin": np.ones((1, 3, 3), dtype=np.float64), + } + model.change_out_bias([sample], preset_out_bias={"energy": [1.0, 2.0, 3.0]}) + np.testing.assert_allclose( + to_numpy_array(model.get_out_bias()).reshape(-1), [1.0, 2.0, 3.0, 0.0, 0.0, 0.0] + ) + + def test_lower_spin_expansion_preserves_virtual_types(model) -> None: """Local and ghost placeholders remain negative in the switched layout.""" extended_coord = np.arange(12, dtype=np.float64).reshape(1, 4, 3) diff --git a/source/tests/common/test_preset_out_bias.py b/source/tests/common/test_preset_out_bias.py index 02cadb51bc..c38e2a3c2c 100644 --- a/source/tests/common/test_preset_out_bias.py +++ b/source/tests/common/test_preset_out_bias.py @@ -18,11 +18,47 @@ normalize_preset_out_bias, override_assigned_bias, preset_out_bias_rows, + recover_nonfinite_preset_bias, remap_preset_out_bias, resolve_preset_out_bias_tables, ) +class TestRecoverNonfinitePresetBias(unittest.TestCase): + def test_recovers_only_assigned_nonfinite_cells(self) -> None: + stored = np.array( + [ + [[np.nan, 0, 0], [7, 0, 0], [np.inf, 0, 0]], + [[1, np.nan, 3], [4, np.inf, 6], [7, 8, -np.inf]], + [[np.nan, 0, 0], [1, 0, 0], [2, 0, 0]], + ] + ) + original = stored.copy() + preset = { + "energy": [[-10], [-13.6], [-3]], + "tensor": [[10, 11, 12], None, [30, 31, 32]], + } + rows = recover_nonfinite_preset_bias( + preset, ["O", "H", "B"], stored, ["energy", "tensor", "other"], [1, 3, 1] + ) + np.testing.assert_array_equal(rows["energy"], [[-10], [7], [-3]]) + np.testing.assert_array_equal( + rows["tensor"], [[1, 11, 3], [4, np.inf, 6], [7, 8, 32]] + ) + self.assertNotIn("other", rows) + np.testing.assert_array_equal(stored, original) + + def test_finite_adjusted_bias_needs_no_recovery(self) -> None: + rows = recover_nonfinite_preset_bias( + {"energy": [[-10], [-13.6]]}, + ["O", "H"], + np.array([[[7.0], [8.0]]]), + ["energy"], + [1], + ) + self.assertEqual(rows, {}) + + class TestNormalizePresetOutBias(unittest.TestCase): def setUp(self) -> None: self.type_map = ["O", "H", "B"] diff --git a/source/tests/pd/model/test_atomic_model_global_stat.py b/source/tests/pd/model/test_atomic_model_global_stat.py index bacec7e682..023bafb277 100644 --- a/source/tests/pd/model/test_atomic_model_global_stat.py +++ b/source/tests/pd/model/test_atomic_model_global_stat.py @@ -4,6 +4,9 @@ from pathlib import ( Path, ) +from unittest.mock import ( + patch, +) import h5py import numpy as np @@ -360,6 +363,42 @@ def cvt_ret(x): out_bias, _ = md0._fetch_out_stat(["foo", "bar"]) np.testing.assert_almost_equal(to_numpy_array(out_bias["foo"]), foo_bias) np.testing.assert_almost_equal(to_numpy_array(out_bias["bar"]), bar_bias) + + # Standalone change-bias explicitly ignores the checkpoint preset. + adjusted_bias = foo_bias + 7.0 + md0._store_out_stat({"foo": to_paddle_tensor(adjusted_bias)}) + forward = md0._get_forward_wrapper_func() + seen_biases = [] + + def check_forward(*args, **kwargs): + bias, _ = md0._fetch_out_stat(["foo"]) + seen_biases.append(to_numpy_array(bias["foo"]).copy()) + return forward(*args, **kwargs) + + with patch.object(md0, "_get_forward_wrapper_func", return_value=check_forward): + BaseAtomicModel.change_out_bias( + md0, + self.merged_output_stat, + bias_adjust_mode="change-by-statistic", + preset_out_bias={}, + ) + self.assertTrue(seen_biases) + for seen_bias in seen_biases: + np.testing.assert_almost_equal(seen_bias, adjusted_bias) + out_bias, _ = md0._fetch_out_stat(["foo", "bar"]) + np.testing.assert_almost_equal( + to_numpy_array(out_bias["foo"]), np.array([2.0, -5.0]).reshape(2, 1) + ) + assert not np.allclose(to_numpy_array(out_bias["bar"]), bar_bias) + + # The override does not remove the model's configured preset. + np.testing.assert_almost_equal(md0.preset_out_bias["foo"], foo_bias) + BaseAtomicModel.change_out_bias( + md0, self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + out_bias, _ = md0._fetch_out_stat(["foo", "bar"]) + np.testing.assert_almost_equal(to_numpy_array(out_bias["foo"]), foo_bias) + np.testing.assert_almost_equal(to_numpy_array(out_bias["bar"]), bar_bias) # 4. a preset that leaves an observed type unassigned is rejected md1 = DPAtomicModel( ds, diff --git a/source/tests/pd/test_change_bias.py b/source/tests/pd/test_change_bias.py index 2d87b739ff..36c2f0290c 100644 --- a/source/tests/pd/test_change_bias.py +++ b/source/tests/pd/test_change_bias.py @@ -140,6 +140,29 @@ def test_change_bias_with_user_defined(self): ) np.testing.assert_allclose(updated_bias.numpy(), expected_bias.numpy()) + def test_change_bias_ignores_checkpoint_preset(self): + checkpoint = paddle.load(str(self.model_path)) + preset = {"energy": [1.0e6] * len(self.config["model"]["type_map"])} + checkpoint["model"]["_extra_state"]["model_params"]["preset_out_bias"] = preset + paddle.save(checkpoint, str(self.model_path)) + run_dp( + f"dp --pd change-bias {self.model_path!s} -s {self.data_file[0]} -o {self.model_path_data_bias!s}" + ) + updated = paddle.load(str(self.model_path_data_bias)) + model_params = updated["model"]["_extra_state"]["model_params"] + self.assertEqual(model_params["preset_out_bias"], preset) + wrapper = ModelWrapper(get_model_for_wrapper(model_params)) + wrapper.set_state_dict(updated["model"]) + expected_model = model_change_out_bias( + self.trainer.wrapper.model["Default"], + self.sampled, + _bias_adjust_mode="change-by-statistic", + ) + np.testing.assert_allclose( + wrapper.model["Default"].get_out_bias().numpy(), + expected_model.get_out_bias().numpy(), + ) + def tearDown(self): for f in os.listdir("."): if f.startswith("change-bias-model") and f.endswith(".pd"): diff --git a/source/tests/pt/model/test_atomic_model_global_stat.py b/source/tests/pt/model/test_atomic_model_global_stat.py index a4bca815b9..78c877bce9 100644 --- a/source/tests/pt/model/test_atomic_model_global_stat.py +++ b/source/tests/pt/model/test_atomic_model_global_stat.py @@ -7,6 +7,9 @@ from typing import ( NoReturn, ) +from unittest.mock import ( + patch, +) import h5py import numpy as np @@ -386,6 +389,42 @@ def cvt_ret(x): out_bias, _ = md0._fetch_out_stat(["foo", "bar"]) np.testing.assert_almost_equal(to_numpy_array(out_bias["foo"]), foo_bias) np.testing.assert_almost_equal(to_numpy_array(out_bias["bar"]), bar_bias) + + # Standalone change-bias explicitly ignores the checkpoint preset. + adjusted_bias = foo_bias + 7.0 + md0._store_out_stat({"foo": to_torch_tensor(adjusted_bias)}) + forward = md0._get_forward_wrapper_func() + seen_biases = [] + + def check_forward(*args, **kwargs): + bias, _ = md0._fetch_out_stat(["foo"]) + seen_biases.append(to_numpy_array(bias["foo"]).copy()) + return forward(*args, **kwargs) + + with patch.object(md0, "_get_forward_wrapper_func", return_value=check_forward): + BaseAtomicModel.change_out_bias( + md0, + self.merged_output_stat, + bias_adjust_mode="change-by-statistic", + preset_out_bias={}, + ) + self.assertTrue(seen_biases) + for seen_bias in seen_biases: + np.testing.assert_almost_equal(seen_bias, adjusted_bias) + out_bias, _ = md0._fetch_out_stat(["foo", "bar"]) + np.testing.assert_almost_equal( + to_numpy_array(out_bias["foo"]), np.array([2.0, -5.0]).reshape(2, 1) + ) + assert not np.allclose(to_numpy_array(out_bias["bar"]), bar_bias) + + # The override does not remove the model's configured preset. + np.testing.assert_almost_equal(md0.preset_out_bias["foo"], foo_bias) + BaseAtomicModel.change_out_bias( + md0, self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + out_bias, _ = md0._fetch_out_stat(["foo", "bar"]) + np.testing.assert_almost_equal(to_numpy_array(out_bias["foo"]), foo_bias) + np.testing.assert_almost_equal(to_numpy_array(out_bias["bar"]), bar_bias) # 4. a preset that leaves an observed type unassigned is rejected md1 = DPAtomicModel( ds, @@ -659,6 +698,43 @@ def test_nonfinite_assigned_bias_from_checkpoint(self) -> None: to_numpy_array(loaded.get_out_bias()).reshape(-1), bias, atol=1e-10 ) + def test_empty_override_recovers_nonfinite_checkpoint_bias(self) -> None: + for value in (np.nan, np.inf, -np.inf): + with self.subTest(value=value): + loaded = EnergyModel.deserialize(self.model.serialize()) + loaded.atomic_model.preset_out_bias = { + "energy": [[-10.0], [-13.6], [-3.0]] + } + loaded.atomic_model.out_bias[0, :, 0] = to_torch_tensor( + np.array([value, 7.0, value]) + ) + loaded = EnergyModel.deserialize(loaded.serialize()) + forward = loaded.predict_atomic_outputs_for_stat + seen = [] + + def check_forward(*args, **kwargs): + seen.append( + to_numpy_array(loaded.get_out_bias()).reshape(-1).copy() + ) + return forward(*args, **kwargs) + + with patch.object( + loaded, "predict_atomic_outputs_for_stat", check_forward + ): + loaded.change_out_bias(self.sampled, preset_out_bias={}) + self.assertTrue(seen) + for baseline in seen: + np.testing.assert_allclose(baseline, [-10.0, 7.0, -3.0]) + bias = to_numpy_array(loaded.get_out_bias()).reshape(-1) + self.assertTrue(np.isfinite(bias).all()) + self.assertAlmostEqual(bias[2], -3.0) + prediction = loaded(self.sampled[0]["coord"], self.sampled[0]["atype"]) + np.testing.assert_allclose( + to_numpy_array(prediction["energy"]), + to_numpy_array(self.sampled[0]["energy"]), + atol=1e-5, + ) + def test_change_type_map_remaps_preset(self) -> None: self.model.change_out_bias(self.sampled, bias_adjust_mode="set-by-statistic") # H is dropped, O keeps its preset, B and C are unassigned diff --git a/source/tests/pt/model/test_sezm_spin_model.py b/source/tests/pt/model/test_sezm_spin_model.py index ffb86609fb..cf8e5b20f1 100644 --- a/source/tests/pt/model/test_sezm_spin_model.py +++ b/source/tests/pt/model/test_sezm_spin_model.py @@ -376,6 +376,21 @@ def test_bridging_is_rejected(self) -> None: with self.assertRaisesRegex(ValueError, "native"): get_model(self._build_model_params(bridging_method="ZBL")) + def test_real_type_sequence_bias_override(self) -> None: + model = get_model(self._build_model_params()).to(self.device) + sample = { + "coord": self.coord, + "atype": self.atype, + "spin": self.spin, + "box": self.box, + } + model.change_out_bias([sample], preset_out_bias={"energy": [1.0, 2.0]}) + bias = model.get_out_bias() + torch.testing.assert_close( + bias.reshape(-1), + torch.tensor([1.0, 2.0, 0.0, 0.0], dtype=bias.dtype, device=bias.device), + ) + @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) def test_compile_matches_eager(self) -> None: """Compiled SeZM spin path should match eager predictions.""" diff --git a/source/tests/pt/model/test_spin_model_virtual_types.py b/source/tests/pt/model/test_spin_model_virtual_types.py index b8d8dbf772..8193751b06 100644 --- a/source/tests/pt/model/test_spin_model_virtual_types.py +++ b/source/tests/pt/model/test_spin_model_virtual_types.py @@ -61,6 +61,19 @@ def _tensor(array: np.ndarray) -> torch.Tensor: class TestPtSpinModelVirtualTypes(unittest.TestCase): """Every pt lookup must mask ``atype < 0`` exactly as dpmodel does.""" + def test_real_type_sequence_bias_override(self) -> None: + model = get_pt_model(MODEL_CONFIG) + sample = { + "coord": _tensor(np.arange(9, dtype=np.float64).reshape(1, 3, 3)), + "atype": _tensor(np.array([[0, 1, 2]], dtype=np.int64)), + "spin": _tensor(np.ones((1, 3, 3), dtype=np.float64)), + } + model.change_out_bias([sample], preset_out_bias={"energy": [1.0, 2.0, 3.0]}) + np.testing.assert_allclose( + to_numpy_array(model.get_out_bias()).reshape(-1), + [1.0, 2.0, 3.0, 0.0, 0.0, 0.0], + ) + def test_dense_spin_expansion_preserves_virtual_types(self) -> None: model = get_pt_model(MODEL_CONFIG) coord = np.arange(9, dtype=np.float64).reshape(1, 3, 3) diff --git a/source/tests/pt/test_change_bias.py b/source/tests/pt/test_change_bias.py index 746656ac8c..f7e3060293 100644 --- a/source/tests/pt/test_change_bias.py +++ b/source/tests/pt/test_change_bias.py @@ -144,6 +144,32 @@ def test_change_bias_with_data_sys_file(self) -> None: expected_bias = expected_model.get_out_bias() torch.testing.assert_close(updated_bias, expected_bias) + def test_change_bias_ignores_checkpoint_preset(self) -> None: + checkpoint = torch.load( + str(self.model_path), map_location=DEVICE, weights_only=True + ) + preset = {"energy": [1.0e6] * len(self.config["model"]["type_map"])} + checkpoint["model"]["_extra_state"]["model_params"]["preset_out_bias"] = preset + torch.save(checkpoint, self.model_path) + run_dp( + f"dp --pt change-bias {self.model_path!s} -s {self.data_file[0]} -o {self.model_path_data_bias!s}" + ) + updated = torch.load( + str(self.model_path_data_bias), map_location=DEVICE, weights_only=True + ) + model_params = updated["model"]["_extra_state"]["model_params"] + self.assertEqual(model_params["preset_out_bias"], preset) + wrapper = ModelWrapper(get_model_for_wrapper(model_params)) + wrapper.load_state_dict(updated["model"]) + expected_model = model_change_out_bias( + self.trainer.wrapper.model["Default"], + self.sampled, + _bias_adjust_mode="change-by-statistic", + ) + torch.testing.assert_close( + wrapper.model["Default"].get_out_bias(), expected_model.get_out_bias() + ) + def test_change_bias_with_user_defined(self) -> None: user_bias = [0.1, 3.2, -0.5] run_dp( diff --git a/source/tests/pt_expt/test_change_bias.py b/source/tests/pt_expt/test_change_bias.py index 2401aa59c5..719ada14b3 100644 --- a/source/tests/pt_expt/test_change_bias.py +++ b/source/tests/pt_expt/test_change_bias.py @@ -19,18 +19,27 @@ from deepmd.pt_expt.model.get_model import ( get_model, ) +from deepmd.pt_expt.train.training import ( + model_change_out_bias, +) from deepmd.pt_expt.train.wrapper import ( ModelWrapper, ) from deepmd.pt_expt.utils.env import ( DEVICE, ) +from deepmd.pt_expt.utils.stat import ( + make_stat_input, +) from deepmd.utils.argcheck import ( normalize, ) from deepmd.utils.compat import ( update_deepmd_input, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) def to_numpy(x): @@ -184,6 +193,7 @@ def setUpClass(cls) -> None: # Build & train 1-step model config = _make_config(data_dir) + config["model"]["preset_out_bias"] = {"energy": [-10.0, -13.6]} config = update_deepmd_input(config, warning=False) config = normalize(config) config["training"]["save_ckpt"] = "model.ckpt" @@ -193,8 +203,18 @@ def setUpClass(cls) -> None: cls.model_path = os.path.join(cls.tmpdir, "model.ckpt.pt") - # Record original bias - cls.original_bias = to_numpy(trainer.wrapper.model["Default"].get_out_bias()) + # A checkpoint may already carry a downstream-adjusted finite bias. + model = trainer.wrapper.model["Default"] + old_bias = model.get_out_bias() + model.set_out_bias( + torch.as_tensor( + [7.0, 8.0], dtype=old_bias.dtype, device=old_bias.device + ).reshape(old_bias.shape) + ) + checkpoint = torch.load(cls.model_path, map_location=DEVICE, weights_only=True) + checkpoint["model"] = trainer.wrapper.state_dict() + torch.save(checkpoint, cls.model_path) + cls.original_bias = to_numpy(model.get_out_bias()).copy() # Pre-freeze shared .pte and .pt2 files so individual tests don't # each pay the AOTInductor compilation cost (~82s per .pt2). @@ -231,6 +251,44 @@ def _load_model_from_ckpt(self, ckpt_path: str): wrapper.load_state_dict(model_state) return model + def _expected_fitted_bias(self): + model = self._load_model_from_ckpt(self.model_path) + # Compare against the established no-preset statistics path with the + # same network and adjusted bias, rather than starting from the preset. + model.atomic_model.preset_out_bias = None + data = DeepmdDataSystem( + systems=self.data_file, + batch_size=1, + test_size=1, + rcut=model.get_rcut(), + type_map=model.get_type_map(), + ) + from deepmd.pt_expt.train.training import ( + get_loss, + ) + + loss = get_loss({"inference": True}, 1.0, len(model.get_type_map()), model) + data.add_data_requirements(loss.label_requirement) + sampled = make_stat_input(data, min(data.get_nbatches())) + return to_numpy(model_change_out_bias(model, sampled).get_out_bias()) + + def _assert_preset_bias_update(self, model) -> None: + np.testing.assert_allclose( + to_numpy(model.get_out_bias()), self._expected_fitted_bias(), atol=1e-10 + ) + self.assertEqual( + model.atomic_model.preset_out_bias, {"energy": [[-10.0], [-13.6]]} + ) + + def test_default_bias_helper_preserves_preset_behavior(self) -> None: + model = self._load_model_from_ckpt(self.model_path) + np.testing.assert_allclose(to_numpy(model.get_out_bias()), self.original_bias) + sampled = [{"atype": np.array([[0, 1]], dtype=np.int32)}] + model_change_out_bias(model, sampled) + np.testing.assert_allclose( + to_numpy(model.get_out_bias()).reshape(-1), [-10.0, -13.6] + ) + def test_change_bias_with_data(self) -> None: output_path = os.path.join(self.tmpdir, "model_data_bias.pt") run_dp( @@ -239,6 +297,7 @@ def test_change_bias_with_data(self) -> None: ) updated_model = self._load_model_from_ckpt(output_path) updated_bias = to_numpy(updated_model.get_out_bias()) + self._assert_preset_bias_update(updated_model) original_bias = np.array(self.original_bias) # Bias should have changed from the original self.assertFalse( @@ -333,6 +392,11 @@ def test_change_bias_frozen_pte(self) -> None: updated_data = serialize_from_file(output_pte) updated_model = BaseModel.deserialize(updated_data["model"]) updated_bias = to_numpy(updated_model.get_out_bias()) + self._assert_preset_bias_update(updated_model) + self.assertEqual( + updated_data["model_def_script"]["preset_out_bias"], + self.config["model"]["preset_out_bias"], + ) # Bias should have changed self.assertFalse( @@ -362,6 +426,11 @@ def test_change_bias_frozen_pt2(self) -> None: updated_data = serialize_from_file(output_pt2) updated_model = BaseModel.deserialize(updated_data["model"]) updated_bias = to_numpy(updated_model.get_out_bias()) + self._assert_preset_bias_update(updated_model) + self.assertEqual( + updated_data["model_def_script"]["preset_out_bias"], + self.config["model"]["preset_out_bias"], + ) self.assertFalse( np.allclose(original_bias, updated_bias),