Merge pull request #859 from deepmodeling/zjgemi

fix: optimize memory usage in HDF5Dataset
deepmodeling · Sep 11, 2024 · 0f2bcf5 · 0f2bcf5
2 parents 7edb39a + 3fa4046
commit 0f2bcf5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/src/dflow/python/opio.py b/src/dflow/python/opio.py
@@ -22,8 +22,13 @@ class NestedDictPath(NestedDictBase):
 
 
 class HDF5Dataset:
-    def __init__(self, dataset):
-        self.dataset = dataset
+    def __init__(self, file, key):
+        self.file = file
+        self.key = key
+
+    @property
+    def dataset(self):
+        return self.file[self.key]
 
     def __deepcopy__(self, memo=None):
         return self

diff --git a/src/dflow/python/utils.py b/src/dflow/python/utils.py
@@ -101,7 +101,7 @@ def handle_input_artifact(name, sign, slices=None, data_root="/tmp",
         res = None
         for path in path_object:
             f = h5py.File(path, "r")
-            datasets = {k: HDF5Dataset(f[k]) for k in f.keys()}
+            datasets = {k: HDF5Dataset(f, k) for k in f.keys()}
             datasets = expand(datasets)
             if isinstance(datasets, list):
                 if res is None: