use ThreadPoolExecutor to eliminate CUDA re-initialization in data modification during training.

Author: ChiahsinChu

deepmd/pd/utils/dataset.py

@@ -5,6 +5,9 @@
     Dataset,
 )
 
+from deepmd.pd.utils.env import (
+    NUM_WORKERS,
+)
 from deepmd.utils.data import (
     DataRequirementItem,
     DeepmdData,
@@ -32,7 +35,7 @@ def __len__(self):
 
     def __getitem__(self, index):
         """Get a frame from the selected system."""
-        b_data = self._data_system.get_item_paddle(index)
+        b_data = self._data_system.get_item_paddle(index, NUM_WORKERS)
         b_data["natoms"] = self._natoms_vec
         return b_data
 

deepmd/pt/train/training.py

@@ -339,9 +339,9 @@ def get_lr(lr_params: dict[str, Any]) -> LearningRateExp:
             if validation_data is not None:
                 validation_data.add_data_requirement(data_requirement)
             # Preload and apply modifiers to all data before computing statistics
-            training_data.preload_and_modify_all_data()
+            training_data.preload_and_modify_all_data_torch()
             if validation_data is not None:
-                validation_data.preload_and_modify_all_data()
+                validation_data.preload_and_modify_all_data_torch()
             self.get_sample_func = single_model_stat(
                 self.model,
                 model_params.get("data_stat_nbatch", 10),
@@ -385,9 +385,9 @@ def get_lr(lr_params: dict[str, Any]) -> LearningRateExp:
                 if validation_data[model_key] is not None:
                     validation_data[model_key].add_data_requirement(data_requirement)
                 # Preload and apply modifiers to all data before computing statistics
-                training_data[model_key].preload_and_modify_all_data()
+                training_data[model_key].preload_and_modify_all_data_torch()
                 if validation_data[model_key] is not None:
-                    validation_data[model_key].preload_and_modify_all_data()
+                    validation_data[model_key].preload_and_modify_all_data_torch()
                 self.get_sample_func[model_key] = single_model_stat(
                     self.model[model_key],
                     model_params["model_dict"][model_key].get("data_stat_nbatch", 10),

deepmd/pt/utils/dataloader.py

@@ -238,9 +238,9 @@ def print_summary(
                 [ss._data_system.pbc for ss in self.systems],
             )
 
-    def preload_and_modify_all_data(self) -> None:
+    def preload_and_modify_all_data_torch(self) -> None:
         for system in self.systems:
-            system.preload_and_modify_all_data()
+            system.preload_and_modify_all_data_torch()
 
 
 def collate_batch(batch: list[dict[str, Any]]) -> dict[str, Any]:

deepmd/pt/utils/dataset.py

@@ -12,6 +12,9 @@
 from deepmd.pt.modifier import (
     BaseModifier,
 )
+from deepmd.pt.utils.env import (
+    NUM_WORKERS,
+)
 from deepmd.utils.data import (
     DataRequirementItem,
     DeepmdData,
@@ -48,7 +51,7 @@ def __len__(self) -> int:
 
     def __getitem__(self, index: int) -> dict[str, Any]:
         """Get a frame from the selected system."""
-        b_data = self._data_system.get_item_torch(index)
+        b_data = self._data_system.get_item_torch(index, NUM_WORKERS)
         b_data["natoms"] = self._natoms_vec
         return b_data
 
@@ -68,5 +71,5 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N
                 output_natoms_for_type_sel=data_item["output_natoms_for_type_sel"],
             )
 
-    def preload_and_modify_all_data(self) -> None:
-        self._data_system.preload_and_modify_all_data()
+    def preload_and_modify_all_data_torch(self) -> None:
+        self._data_system.preload_and_modify_all_data_torch(NUM_WORKERS)

deepmd/utils/data.py

@@ -253,26 +253,38 @@ def check_test_size(self, test_size: int) -> bool:
         """Check if the system can get a test dataset with `test_size` frames."""
         return self.check_batch_size(test_size)
 
-    def get_item_torch(self, index: int) -> dict:
+    def get_item_torch(
+        self,
+        index: int,
+        num_worker: int,
+    ) -> dict:
         """Get a single frame data . The frame is picked from the data system by index. The index is coded across all the sets.
 
         Parameters
         ----------
         index
             index of the frame
+        num_worker
+            number of workers for parallel data modification
         """
-        return self.get_single_frame(index)
+        return self.get_single_frame(index, num_worker)
 
-    def get_item_paddle(self, index: int) -> dict:
+    def get_item_paddle(
+        self,
+        index: int,
+        num_worker: int,
+    ) -> dict:
         """Get a single frame data . The frame is picked from the data system by index. The index is coded across all the sets.
         Same with PyTorch backend.
 
         Parameters
         ----------
         index
             index of the frame
+        num_worker
+            number of workers for parallel data modification
         """
-        return self.get_single_frame(index)
+        return self.get_single_frame(index, num_worker)
 
     def get_batch(self, batch_size: int) -> dict:
         """Get a batch of data with `batch_size` frames. The frames are randomly picked from the data system.
@@ -383,7 +395,7 @@ def get_natoms_vec(self, ntypes: int) -> np.ndarray:
         tmp = np.append(tmp, natoms_vec)
         return tmp.astype(np.int32)
 
-    def get_single_frame(self, index: int) -> dict:
+    def get_single_frame(self, index: int, num_worker: int) -> dict:
         """Orchestrates loading a single frame efficiently using memmap."""
         # Check if we have a cached modified frame and use_modifier_cache is True
         if (
@@ -488,14 +500,19 @@ def get_single_frame(self, index: int) -> dict:
         frame_data["fid"] = index
 
         if self.modifier is not None:
-            # Apply modifier if it exists
-            self.modifier.modify_data(frame_data, self)
+            with ThreadPoolExecutor(max_workers=num_worker) as executor:
+                # Apply modifier if it exists
+                executor.submit(
+                    self.modifier.modify_data,
+                    frame_data,
+                    self,
+                )
             if self.use_modifier_cache:
                 # Cache the modified frame to avoid recomputation
                 self._modified_frame_cache[index] = copy.deepcopy(frame_data)
         return frame_data
 
-    def preload_and_modify_all_data(self) -> None:
+    def preload_and_modify_all_data_torch(self, num_worker: int) -> None:
         """Preload all frames and apply modifier to cache them.
 
         This method is useful when use_modifier_cache is True and you want to
@@ -507,7 +524,7 @@ def preload_and_modify_all_data(self) -> None:
         log.info("Preloading and modifying all data frames...")
         for i in range(self.nframes):
             if i not in self._modified_frame_cache:
-                self.get_single_frame(i)
+                self.get_single_frame(i, num_worker)
                 if (i + 1) % 100 == 0:
                     log.info(f"Processed {i + 1}/{self.nframes} frames")
         log.info("All frames preloaded and modified.")

source/tests/pt/test_data_modifier.py

@@ -325,7 +325,7 @@ def test_inference(self):
             "type": "scaling_tester",
             "model_name": "frozen_model_dm.pth",
             "sfactor": sfactor,
-            "use_cache": True,
+            "use_cache": self.param[2],
         }
 
         trainer = get_trainer(tmp_config)