feat(tf): implement change-bias command

deepmd/main.py

@@ -721,12 +721,13 @@ def main_parser() -> argparse.ArgumentParser:
     parser_change_bias = subparsers.add_parser(
         "change-bias",
         parents=[parser_log],
-        help="(Supported backend: PyTorch) Change model out bias according to the input data.",
+        help="Change model out bias according to the input data.",
         formatter_class=RawTextArgumentDefaultsHelpFormatter,
         epilog=textwrap.dedent(
             """\
         examples:
             dp change-bias model.pt -s data -n 10 -m change
+            dp --tf change-bias checkpoint_dir -s data -n 10 -m change
         """
         ),
     )

deepmd/tf/entrypoints/__init__.py

@@ -4,6 +4,9 @@
 from ..infer.model_devi import (
     make_model_devi,
 )
+from .change_bias import (
+    change_bias,
+)
 from .compress import (
     compress,
 )
@@ -34,6 +37,7 @@
 )
 
 __all__ = [
+    "change_bias",
     "compress",
     "convert",
     "doc_train_input",

deepmd/tf/entrypoints/change_bias.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""DeePMD change bias entrypoint script."""
+
+import logging
+import os
+import shutil
+from pathlib import (
+    Path,
+)
+from typing import (
+    Optional,
+)
+
+from deepmd.common import (
+    expand_sys_str,
+    j_loader,
+)
+from deepmd.tf.train.run_options import (
+    RunOptions,
+)
+from deepmd.tf.train.trainer import (
+    DPTrainer,
+)
+from deepmd.tf.utils.argcheck import (
+    normalize,
+)
+from deepmd.tf.utils.compat import (
+    update_deepmd_input,
+)
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
+
+__all__ = ["change_bias"]
+
+log = logging.getLogger(__name__)
+
+
+def change_bias(
+    INPUT: str,
+    mode: str = "change",
+    bias_value: Optional[list] = None,
+    datafile: Optional[str] = None,
+    system: str = ".",
+    numb_batch: int = 0,
+    model_branch: Optional[str] = None,
+    output: Optional[str] = None,
+    **kwargs,
+) -> None:
+    """Change model out bias according to the input data.
+
+    Parameters
+    ----------
+    INPUT : str
+        The input checkpoint folder or frozen model file
+    mode : str, optional
+        The mode for changing energy bias, by default "change"
+    bias_value : Optional[list], optional
+        The user defined value for each type, by default None
+    datafile : Optional[str], optional
+        The path to the datafile, by default None
+    system : str, optional
+        The system dir, by default "."
+    numb_batch : int, optional
+        The number of frames for bias changing, by default 0
+    model_branch : Optional[str], optional
+        Model branch chosen for changing bias if multi-task model, by default None
+    output : Optional[str], optional
+        The model after changing bias, by default None
+    """
+    input_path = Path(INPUT)
+
+    # Check if input is a checkpoint directory or frozen model
+    if input_path.is_dir():
+        # Checkpoint directory
+        checkpoint_folder = str(input_path)
+        # Check for valid checkpoint early
+        if not (input_path / "checkpoint").exists():
+            raise RuntimeError(f"No valid checkpoint found in {checkpoint_folder}")
+    elif INPUT.endswith((".pb", ".pbtxt")):
+        # Frozen model - for now, not supported
+        raise NotImplementedError(
+            "Bias changing for frozen models (.pb/.pbtxt) is not yet implemented. "
+            "Please provide a checkpoint directory instead. "
+            "You can train a model to create checkpoints, then use this command "
+            "to modify the bias, and finally freeze the modified model."
+        )
+    else:
+        raise RuntimeError(
+            "The model provided must be a checkpoint directory or frozen model file (.pb/.pbtxt)"
+        )
+
+    bias_adjust_mode = "change-by-statistic" if mode == "change" else "set-by-statistic"
+
+    if bias_value is not None:
+        raise NotImplementedError(
+            "User-defined bias setting is not yet implemented for TensorFlow models. "
+            "Please use the data-based bias adjustment mode."
+        )
+
+    # Load data systems for bias calculation
+    if datafile is not None:
+        with open(datafile) as datalist:
+            all_sys = datalist.read().splitlines()
+    else:
+        all_sys = expand_sys_str(system)
+
+    # Load the data systems
+    data = DeepmdDataSystem(
+        systems=all_sys,
+        batch_size=1,
+        test_size=1,
+        rcut=None,
+        set_prefix="set",
+    )
+
+    # Read the checkpoint to get the model configuration
+    checkpoint_path = Path(checkpoint_folder)
+
+    # Find the input.json file or create a minimal config
+    # We need this to reconstruct the model
+    input_json_path = checkpoint_path / "input.json"
+    if not input_json_path.exists():
+        # Look for input.json in parent directories or common locations
+        for parent in checkpoint_path.parents:
+            potential_input = parent / "input.json"
+            if potential_input.exists():
+                input_json_path = potential_input
+                break
+        else:
+            raise RuntimeError(
+                f"Cannot find input.json configuration file needed to load the model. "
+                f"Please ensure input.json is available in {checkpoint_folder} or its parent directories."
+            )
+
+    # Load the configuration
+    jdata = j_loader(input_json_path)
+
+    # Update and normalize the configuration
+    jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json")
+    jdata = normalize(jdata)
+
+    # Determine output path
+    if output is None:
+        output = str(checkpoint_path) + "_bias_updated"
+
+    # Create trainer to access model methods
+    run_opt = RunOptions(
+        init_model=checkpoint_folder,
+        restart=None,
+        finetune=None,
+        init_frz_model=None,
+        train_data=all_sys,
+        valid_data=None,
+    )
+
+    trainer = DPTrainer(jdata, run_opt)
+
+    # Get the type map from the model
+    type_map = data.get_type_map()
+    if len(type_map) == 0:
+        # If data doesn't have type_map, get from model
+        type_map = trainer.model.get_type_map()
+
+    log.info(f"Changing bias for model with type_map: {type_map}")
+    log.info(f"Using bias adjustment mode: {bias_adjust_mode}")
+
+    # Use the trainer's change energy bias functionality
+    trainer._change_energy_bias(
+        data,
+        checkpoint_folder,  # Use checkpoint as frozen model path for compatibility
+        type_map,
+        bias_adjust_mode=bias_adjust_mode,
+    )
+
+    # Save the updated model
+    shutil.copytree(checkpoint_folder, output, dirs_exist_ok=True)
+    trainer.save_checkpoint(os.path.join(output, "model.ckpt"))
+
+    log.info(f"Bias changing complete. Updated model saved to {output}")
+    log.info(
+        f"You can now freeze this model using: dp freeze -c {output} -o model_updated.pb"
+    )

deepmd/tf/entrypoints/main.py

@@ -2,6 +2,7 @@
 """DeePMD-Kit entry point module."""
 
 import argparse
+import logging
 from pathlib import (
     Path,
 )
@@ -22,6 +23,7 @@
     clear_session,
 )
 from deepmd.tf.entrypoints import (
+    change_bias,
     compress,
     convert,
     freeze,
@@ -37,6 +39,8 @@
 
 __all__ = ["get_ll", "main", "main_parser", "parse_args"]
 
+log = logging.getLogger(__name__)
+
 
 def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None:
     """DeePMD-Kit entry point.
@@ -86,6 +90,8 @@
         compress(**dict_args)
     elif args.command == "convert-from":
         convert(**dict_args)
+    elif args.command == "change-bias":
+        change_bias(**dict_args)
     elif args.command == "train-nvnmd":  # nvnmd
         train_nvnmd(**dict_args)
     elif args.command is None:

source/tests/tf/test_change_bias.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import tempfile
+import unittest
+from pathlib import (
+    Path,
+)
+
+from deepmd.tf.entrypoints.change_bias import (
+    change_bias,
+)
+
+
+class TestChangeBias(unittest.TestCase):
+    def setUp(self):
+        """Set up test fixtures."""
+        self.temp_dir = tempfile.mkdtemp()
+        self.temp_path = Path(self.temp_dir)
+
+    def tearDown(self):
+        """Clean up test fixtures."""
+        import shutil
+
+        shutil.rmtree(self.temp_dir, ignore_errors=True)
+
+    def test_change_bias_frozen_model_not_implemented(self):
+        """Test that frozen model support raises NotImplementedError."""
+        fake_pb = self.temp_path / "model.pb"
+        fake_pb.write_text("fake model content")
+
+        with self.assertRaises(NotImplementedError) as cm:
+            change_bias(
+                input_file=str(fake_pb),
+                mode="change",
+                system=".",
+            )
+
+        self.assertIn("Bias changing for frozen models", str(cm.exception))
+        self.assertIn(".pb/.pbtxt", str(cm.exception))
+
+    def test_change_bias_invalid_model_type(self):
+        """Test that invalid model types raise RuntimeError."""
+        fake_model = self.temp_path / "model.xyz"
+        fake_model.write_text("fake model content")
+
+        with self.assertRaises(RuntimeError) as cm:
+            change_bias(
+                input_file=str(fake_model),
+                mode="change",
+                system=".",
+            )
+
+        self.assertIn("checkpoint directory or frozen model file", str(cm.exception))
+
+    def test_change_bias_no_checkpoint_in_directory(self):
+        """Test that missing checkpoint in directory raises RuntimeError."""
+        fake_dir = self.temp_path / "fake_checkpoint"
+        fake_dir.mkdir()
+
+        # Create a fake data system for the test
+        fake_data_dir = self.temp_path / "fake_data"
+        fake_data_dir.mkdir()
+        fake_set_dir = fake_data_dir / "set.000"
+        fake_set_dir.mkdir()
+
+        with self.assertRaises(RuntimeError) as cm:
+            change_bias(
+                input_file=str(fake_dir),
+                mode="change",
+                system=str(fake_data_dir),
+            )
+
+        self.assertIn("No valid checkpoint found", str(cm.exception))
+
+    def test_change_bias_user_defined_not_implemented(self):
+        """Test that user-defined bias raises NotImplementedError."""
+        fake_dir = self.temp_path / "fake_checkpoint"
+        fake_dir.mkdir()
+        (fake_dir / "checkpoint").write_text("fake checkpoint")
+
+        with self.assertRaises(NotImplementedError) as cm:
+            change_bias(
+                input_file=str(fake_dir),
+                mode="change",
+                bias_value=[1.0, 2.0],
+                system=".",
+            )
+
+        self.assertIn(
+            "User-defined bias setting is not yet implemented", str(cm.exception)
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()