feat(pt_expt): add dp change-bias support

deepmd/pt_expt/entrypoints/main.py

@@ -218,6 +218,177 @@ def freeze(
     log.info("Saved frozen model to %s", output)
 
 
+def change_bias(
+    input_file: str,
+    mode: str = "change",
+    bias_value: list | None = None,
+    datafile: str | None = None,
+    system: str = ".",
+    numb_batch: int = 0,
+    model_branch: str | None = None,
+    output: str | None = None,
+) -> None:
+    """Change the output bias of a pt_expt model.
+
+    Parameters
+    ----------
+    input_file : str
+        Path to the model file (.pt checkpoint or .pte frozen model).
+    mode : str
+        ``"change"`` or ``"set"``.
+    bias_value : list or None
+        User-defined bias values (one per type).
+    datafile : str or None
+        File listing data system paths.
+    system : str
+        Data system path (used when *datafile* is None).
+    numb_batch : int
+        Number of batches for statistics (0 = all).
+    model_branch : str or None
+        Branch name for multi-task models.
+    output : str or None
+        Output file path.
+    """
+    import torch
+
+    from deepmd.common import (
+        expand_sys_str,
+    )
+    from deepmd.dpmodel.common import (
+        to_numpy_array,
+    )
+    from deepmd.pt_expt.model.get_model import (
+        get_model,
+    )
+    from deepmd.pt_expt.train.training import (
+        get_additional_data_requirement,
+        get_loss,
+        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.serialization import (
+        deserialize_to_file,
+        serialize_from_file,
+    )
+    from deepmd.pt_expt.utils.stat import (
+        make_stat_input,
+    )
+
+    if input_file.endswith(".pt"):
+        old_state_dict = torch.load(input_file, map_location=DEVICE, weights_only=True)
+        if "model" in old_state_dict:
+            model_state_dict = old_state_dict["model"]
+        else:
+            model_state_dict = old_state_dict
+        extra_state = model_state_dict.get("_extra_state")
+        if not isinstance(extra_state, dict) or "model_params" not in extra_state:
+            raise ValueError(
+                f"Unsupported checkpoint format at '{input_file}': missing "
+                "'_extra_state.model_params' in model state dict."
+            )
+        model_params = extra_state["model_params"]
+    elif input_file.endswith((".pte", ".pt2")):
+        pte_data = serialize_from_file(input_file)
+        from deepmd.pt_expt.model.model import (
+            BaseModel,
+        )
+
+        model_to_change = BaseModel.deserialize(pte_data["model"])
+        model_params = None
+    else:
+        raise RuntimeError(
+            "The model provided must be a checkpoint file with a .pt extension "
+            "or a frozen model with a .pte/.pt2 extension"
+        )
+
+    bias_adjust_mode = "change-by-statistic" if mode == "change" else "set-by-statistic"
+
+    if input_file.endswith(".pt"):
+        multi_task = "model_dict" in model_params
+        if multi_task:
+            raise NotImplementedError(
+                "Multi-task change-bias is not yet supported for the pt_expt backend."
+            )
+        type_map = model_params["type_map"]
+        model = get_model(model_params)
+        wrapper = ModelWrapper(model)
+        wrapper.load_state_dict(model_state_dict)
+        model_to_change = model
+
+    if input_file.endswith((".pte", ".pt2")):
+        type_map = model_to_change.get_type_map()
+
+    if bias_value is not None:
+        assert "energy" in model_to_change.model_output_type(), (
+            "User-defined bias is only available for energy model!"
+        )
+        assert len(bias_value) == len(type_map), (
+            f"The number of elements in the bias should be the same as "
+            f"that in the type_map: {type_map}."
+        )
+        old_bias = model_to_change.get_out_bias()
+        bias_to_set = torch.tensor(
+            bias_value, dtype=old_bias.dtype, device=old_bias.device
+        ).view(old_bias.shape)
+        model_to_change.set_out_bias(bias_to_set)
+        log.info(
+            f"Change output bias of {type_map!s} "
+            f"from {to_numpy_array(old_bias).reshape(-1)!s} "
+            f"to {to_numpy_array(bias_to_set).reshape(-1)!s}."
+        )
+    else:
+        if datafile is not None:
+            with open(datafile) as datalist:
+                all_sys = datalist.read().splitlines()
+        else:
+            all_sys = expand_sys_str(system)
+        data_systems = process_systems(all_sys)
+        data = DeepmdDataSystem(
+            systems=data_systems,
+            batch_size=1,
+            test_size=1,
+            rcut=model_to_change.get_rcut(),
+            type_map=type_map,
+        )
+        mock_loss = get_loss({"inference": True}, 1.0, len(type_map), model_to_change)
+        data.add_data_requirements(mock_loss.label_requirement)
+        data.add_data_requirements(get_additional_data_requirement(model_to_change))
+        nbatches = numb_batch if numb_batch != 0 else max(data.get_nbatches())
+        sampled_data = make_stat_input(data, nbatches)
+        model_to_change = model_change_out_bias(
+            model_to_change, sampled_data, _bias_adjust_mode=bias_adjust_mode
+        )
+
+    if input_file.endswith(".pt"):
+        output_path = (
+            output if output is not None else input_file.replace(".pt", "_updated.pt")
+        )
+        wrapper = ModelWrapper(model_to_change)
+        if "model" in old_state_dict:
+            old_state_dict["model"] = wrapper.state_dict()
+            old_state_dict["model"]["_extra_state"] = extra_state
+        else:
+            old_state_dict = wrapper.state_dict()
+            old_state_dict["_extra_state"] = extra_state
+        torch.save(old_state_dict, output_path)
+    elif input_file.endswith((".pte", ".pt2")):
+        output_path = (
+            output
+            if output is not None
+            else input_file.replace(".pte", "_updated.pte").replace(
+                ".pt2", "_updated.pt2"
+            )
+        )
+        model_dict = model_to_change.serialize()
+        deserialize_to_file(output_path, {"model": model_dict})
+    log.info(f"Saved model to {output_path}")
+
+
 def main(args: list[str] | argparse.Namespace | None = None) -> None:
     """Entry point for the pt_expt backend CLI.
 
@@ -275,6 +446,17 @@ def main(args: list[str] | argparse.Namespace | None = None) -> None:
         if not FLAGS.output.endswith((".pte", ".pt2")):
             FLAGS.output = str(Path(FLAGS.output).with_suffix(".pte"))
         freeze(model=FLAGS.model, output=FLAGS.output, head=FLAGS.head)
+    elif FLAGS.command == "change-bias":
+        change_bias(
+            input_file=FLAGS.INPUT,
+            mode=FLAGS.mode,
+            bias_value=FLAGS.bias_value,
+            datafile=FLAGS.datafile,
+            system=FLAGS.system,
+            numb_batch=FLAGS.numb_batch,
+            model_branch=FLAGS.model_branch,
+            output=FLAGS.output,
+        )
     else:
         raise RuntimeError(
             f"Unsupported command '{FLAGS.command}' for the pt_expt backend."

deepmd/pt_expt/train/training.py

@@ -884,3 +884,27 @@ def print_on_training(
         line += f"   {cur_lr:8.1e}\n"
         fout.write(line)
         fout.flush()
+
+
+def model_change_out_bias(
+    _model: Any,
+    _sample_func: list[dict],
+    _bias_adjust_mode: str = "change-by-statistic",
+) -> Any:
+    old_bias = deepcopy(_model.get_out_bias())
+    _model.change_out_bias(
+        _sample_func,
+        bias_adjust_mode=_bias_adjust_mode,
+    )
+    new_bias = deepcopy(_model.get_out_bias())
+    model_type_map = _model.get_type_map()
+    from deepmd.dpmodel.common import (
+        to_numpy_array,
+    )
+
+    log.info(
+        f"Change output bias of {model_type_map!s} "
+        f"from {to_numpy_array(old_bias).reshape(-1)[: len(model_type_map)]!s} "
+        f"to {to_numpy_array(new_bias).reshape(-1)[: len(model_type_map)]!s}."
+    )
+    return _model