bug fix

Author: OutisLi

deepmd/pt/utils/stat.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import logging
+import os
 from collections import (
     defaultdict,
 )
@@ -60,7 +61,7 @@ def make_stat_input(
 
     lst = []
     # I/O intensive, set a larger number of workers
-    with ThreadPoolExecutor(max_workers=256) as executor:
+    with ThreadPoolExecutor(min(128, (os.cpu_count() or 1) * 6)) as executor:
         lst = list(executor.map(_process_one_dataset, args_list))
     log.info("Finished packing data.")
     return lst
@@ -120,7 +121,7 @@ def _process_one_dataset(args: tuple[Any, int, int]) -> dict[str, Any]:
         if isinstance(sys_stat[key], np.float32):
             pass
         elif isinstance(sys_stat[key], list):
-            if sys_stat[key][0] is None:
+            if len(sys_stat[key]) == 0 or sys_stat[key][0] is None:
                 sys_stat[key] = None
             else:
                 sys_stat[key] = torch.cat(sys_stat[key], dim=0)

deepmd/utils/data.py

@@ -341,8 +341,13 @@ def get_numb_set(self) -> int:
 
     def get_numb_batch(self, batch_size: int, set_idx: int) -> int:
         """Get the number of batches in a set."""
-        # Directly obtain the number of frames to avoid loading the entire dataset
-        nframes = self._get_nframes(self.dirs[set_idx])
+        set_name = self.dirs[set_idx]
+        if isinstance(set_name, DPH5Path):
+            data = self._load_set(set_name)
+            nframes = data["coord"].shape[0]
+        else:
+            # Directly obtain the number of frames to avoid loading the entire dataset
+            nframes = self._get_nframes(set_name)
         ret = nframes // batch_size
         if ret == 0:
             ret = 1
@@ -587,7 +592,7 @@ def _get_nframes(self, set_name: DPPath) -> int:
         if not isinstance(set_name, DPPath):
             set_name = DPPath(set_name)
         path = set_name / "coord.npy"
-        # Read only the header to get shape without creating memmap
+        # Read only the header to get shape
         with open(str(path), "rb") as f:
             version = np.lib.format.read_magic(f)
             if version[0] == 1:

deepmd/utils/env_mat_stat.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import logging
+import os
 from abc import (
     ABC,
     abstractmethod,
@@ -160,7 +161,7 @@ def load_stats(self, path: DPPath) -> None:
         if not files_to_load:
             raise ValueError(f"No statistics files found in {path}.")
 
-        with ThreadPoolExecutor(max_workers=128) as executor:
+        with ThreadPoolExecutor(min(64, (os.cpu_count() or 1) * 4)) as executor:
             results = executor.map(self._load_stat_file, files_to_load)
 
         for name, stat_item in results: