add UT for data modifier in pt model training

Author: ChiahsinChu

deepmd/pt/train/training.py

@@ -391,6 +391,9 @@ def get_lr(lr_params: dict[str, Any]) -> LearningRateExp:
                     if self.finetune_links is not None
                     else False,
                 )
+                training_data[model_key].preload_and_modify_all_data()
+                if validation_data[model_key] is not None:
+                    validation_data[model_key].preload_and_modify_all_data()
                 (
                     self.training_dataloader[model_key],
                     self.training_data[model_key],

deepmd/pt/utils/dataloader.py

@@ -242,6 +242,10 @@ def preload_and_modify_all_data(self) -> None:
         for system in self.systems:
             system.preload_and_modify_all_data()
 
+    # def clear_modified_frame_cache(self) -> None:
+    #     for system in self.systems:
+    #         system.clear_modified_frame_cache()
+
 
 def collate_batch(batch: list[dict[str, Any]]) -> dict[str, Any]:
     example = batch[0]

deepmd/pt/utils/dataset.py

@@ -70,3 +70,6 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N
 
     def preload_and_modify_all_data(self) -> None:
         self._data_system.preload_and_modify_all_data()
+
+    # def clear_modified_frame_cache(self) -> None:
+    #     self._data_system.clear_modified_frame_cache()

deepmd/utils/data.py

@@ -512,14 +512,14 @@ def preload_and_modify_all_data(self) -> None:
                     log.info(f"Processed {i + 1}/{self.nframes} frames")
         log.info("All frames preloaded and modified.")
 
-    def clear_modified_frame_cache(self) -> None:
-        """Clear the modified frame cache.
-
-        This method is useful when you want to free up memory or force
-        recomputation of modified frames.
-        """
-        self._modified_frame_cache.clear()
-        log.info("Modified frame cache cleared.")
+    # def clear_modified_frame_cache(self) -> None:
+    #     """Clear the modified frame cache.
+
+    #     This method is useful when you want to free up memory or force
+    #     recomputation of modified frames.
+    #     """
+    #     self._modified_frame_cache.clear()
+    #     log.info("Modified frame cache cleared.")
 
     def avg(self, key: str) -> float:
         """Return the average value of an item."""

source/tests/pt/test_data_modifier.py

@@ -0,0 +1,232 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import json
+import os
+import unittest
+from pathlib import (
+    Path,
+)
+
+import numpy as np
+
+from deepmd.pt.entrypoints.main import (
+    get_trainer,
+)
+from deepmd.pt.modifier.base_modifier import (
+    BaseModifier,
+)
+from deepmd.pt.utils.utils import (
+    to_numpy_array,
+)
+from deepmd.utils.argcheck import (
+    modifier_args_plugin,
+)
+from deepmd.utils.data import (
+    DeepmdData,
+)
+
+
+@modifier_args_plugin.register("random_tester")
+def modifier_random_tester() -> list:
+    return []
+
+
+@modifier_args_plugin.register("zero_tester")
+def modifier_zero_tester() -> list:
+    return []
+
+
+@BaseModifier.register("random_tester")
+class ModifierRandomTester(BaseModifier):
+    def __new__(cls) -> BaseModifier:
+        return super().__new__(cls)
+
+    def __init__(self) -> None:
+        """Construct a basic model for different tasks."""
+        super().__init__()
+        self.modifier_type = "tester"
+
+    def serialize(self) -> dict:
+        """Serialize the modifier.
+
+        Returns
+        -------
+        dict
+            The serialized data
+        """
+        data = {
+            "@class": "Modifier",
+            "type": self.modifier_type,
+            "@version": 3,
+        }
+        return data
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "BaseModifier":
+        """Deserialize the modifier.
+
+        Parameters
+        ----------
+        data : dict
+            The serialized data
+
+        Returns
+        -------
+        BaseModifier
+            The deserialized modifier
+        """
+        data = data.copy()
+        modifier = cls(**data)
+        return modifier
+
+    def modify_data(self, data: dict, data_sys: DeepmdData) -> None:
+        """Multiply by a random factor."""
+        if (
+            "find_energy" not in data
+            and "find_force" not in data
+            and "find_virial" not in data
+        ):
+            return
+
+        if "find_energy" in data and data["find_energy"] == 1.0:
+            data["energy"] = data["energy"] * np.random.Generator()
+        if "find_force" in data and data["find_force"] == 1.0:
+            data["force"] = data["force"] * np.random.Generator()
+        if "find_virial" in data and data["find_virial"] == 1.0:
+            data["virial"] = data["virial"] * np.random.Generator()
+
+
+@BaseModifier.register("zero_tester")
+class ModifierZeroTester(BaseModifier):
+    def __new__(cls) -> BaseModifier:
+        return super().__new__(cls)
+
+    def __init__(self) -> None:
+        """Construct a basic model for different tasks."""
+        super().__init__()
+        self.modifier_type = "tester"
+
+    def serialize(self) -> dict:
+        """Serialize the modifier.
+
+        Returns
+        -------
+        dict
+            The serialized data
+        """
+        data = {
+            "@class": "Modifier",
+            "type": self.modifier_type,
+            "@version": 3,
+        }
+        return data
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "BaseModifier":
+        """Deserialize the modifier.
+
+        Parameters
+        ----------
+        data : dict
+            The serialized data
+
+        Returns
+        -------
+        BaseModifier
+            The deserialized modifier
+        """
+        data = data.copy()
+        modifier = cls(**data)
+        return modifier
+
+    def modify_data(self, data: dict, data_sys: DeepmdData) -> None:
+        """Multiply by a random factor."""
+        if (
+            "find_energy" not in data
+            and "find_force" not in data
+            and "find_virial" not in data
+        ):
+            return
+
+        if "find_energy" in data and data["find_energy"] == 1.0:
+            data["energy"] -= data["energy"]
+        if "find_force" in data and data["find_force"] == 1.0:
+            data["force"] -= data["force"]
+        if "find_virial" in data and data["find_virial"] == 1.0:
+            data["virial"] -= data["virial"]
+
+
+class TestDataModifier(unittest.TestCase):
+    def setUp(self) -> None:
+        """Set up test fixtures."""
+        input_json = str(Path(__file__).parent / "water/se_e2_a.json")
+        with open(input_json, encoding="utf-8") as f:
+            config = json.load(f)
+        config["training"]["numb_steps"] = 10
+        config["training"]["save_freq"] = 1
+        config["learning_rate"]["start_lr"] = 1.0
+        config["training"]["training_data"]["systems"] = [
+            str(Path(__file__).parent / "water/data/single")
+        ]
+        config["training"]["validation_data"]["systems"] = [
+            str(Path(__file__).parent / "water/data/single")
+        ]
+        self.config = config
+
+    def test_init_modify_data(self):
+        """Ensure modify_data applied."""
+        tmp_config = self.config.copy()
+        # add tester data modifier
+        tmp_config["model"]["modifier"] = {"type": "zero_tester"}
+
+        # data modification is finished in __init__
+        trainer = get_trainer(tmp_config)
+
+        # training data
+        training_data = trainer.get_data(is_train=True)
+        # validation data
+        validation_data = trainer.get_data(is_train=False)
+
+        for dataset in [training_data, validation_data]:
+            for kw in ["energy", "force"]:
+                data = to_numpy_array(dataset[1][kw])
+                np.testing.assert_allclose(data, np.zeros_like(data))
+
+    def test_full_modify_data(self):
+        """Ensure modify_data only applied once."""
+        tmp_config = self.config.copy()
+        # add tester data modifier
+        tmp_config["model"]["modifier"] = {"type": "random_tester"}
+
+        # data modification is finished in __init__
+        trainer = get_trainer(tmp_config)
+
+        # training data
+        training_data_before = trainer.get_data(is_train=True)
+        # validation data
+        validation_data_before = trainer.get_data(is_train=False)
+
+        trainer.run()
+
+        # training data
+        training_data_after = trainer.get_data(is_train=True)
+        # validation data
+        validation_data_after = trainer.get_data(is_train=False)
+
+        for kw in ["energy", "force"]:
+            np.testing.assert_allclose(
+                to_numpy_array(training_data_before[1][kw]),
+                to_numpy_array(training_data_after[1][kw]),
+            )
+            np.testing.assert_allclose(
+                to_numpy_array(validation_data_before[1][kw]),
+                to_numpy_array(validation_data_after[1][kw]),
+            )
+
+    def tearDown(self) -> None:
+        for f in os.listdir("."):
+            if f.startswith("frozen_model") and f.endswith(".pth"):
+                os.remove(f)
+            if f.startswith("model") and f.endswith(".pt"):
+                os.remove(f)
+            if f in ["lcurve.out", "checkpoint"]:
+                os.remove(f)