fix: 修复 dataset 和 checkpoint 的 bug

khaosz/data/checkpoint.py

@@ -70,8 +70,7 @@ class Checkpoint:
             state_dict = torch.load(f)
 
         return cls(
-            optimizer_state_dict=state_dict["optimizer"],
-            scheduler_state_dict=state_dict["scheduler"],
+            state_dict=state_dict,
             epoch=meta["epoch"],
             iteration=meta["iteration"],
             metrics=meta.get("metrics", {}),

khaosz/data/dataset.py

@@ -18,6 +18,9 @@ class BaseSegmentFetcher:
             total += len(seg)
             self.cum_lengths.append(total)
         self.total_length = total if segments else 0
+        
+    def __len__(self) -> int:
+        return self.total_length
 
     def fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
         if not (0 <= begin_idx < self.total_length and 0 <= end_idx <= self.total_length):
@@ -48,6 +51,10 @@ class MultiSegmentFetcher:
             key: BaseSegmentFetcher(segments)
             for key, segments in muti_segments.items()
         }
+    
+    def __len__(self) -> int:
+        len_list = [len(seg) for seg in self.muti_fetchers.values()]
+        return min(len_list)
         
     def key_fetch(self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]) -> Dict:
         fetch_dict = {} 
@@ -73,8 +80,9 @@ class BaseDataset(Dataset, ABC):
         self.total_samples = None
 
     def load(self, load_path: str):
-        self.segments, self.total_samples = load_h5(load_path)
+        self.segments = load_h5(load_path)
         self.fetcher = MultiSegmentFetcher(self.segments)
+        self.total_samples = len(self.fetcher)
         
     def get_index(self, index: int) -> int:
         assert self.total_samples > self.window_size

khaosz/data/file.py

@@ -8,20 +8,17 @@ from torch import Tensor
 from typing import Dict, List, Tuple
 
 
-def save_h5(file_path: str, tensor_group: Dict[str, List[Tensor]]):
-    os.makedirs(os.path.dirname(file_path), exist_ok=True)
-    with h5py.File(file_path, 'w') as f:
+def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
+    os.makedirs(file_path, exist_ok=True)
+    full_file_path = os.path.join(file_path, f"{file_name}.h5")
+    with h5py.File(full_file_path, 'w') as f:
         for key, tensors in tensor_group.items():
             grp = f.create_group(key)
             grp.attrs['num_tensors'] = len(tensors)
             
             for idx, tensor in enumerate(tensors):
                 arr = tensor.cpu().numpy()
-                dset = grp.create_dataset(
-                    f'data_{idx}',
-                    data=arr
-                )
-                dset.attrs['numel'] = tensor.numel()
+                grp.create_dataset(f'data_{idx}', data=arr)
 
 def load_h5(file_path: str) -> Tuple[Dict[str, List[Tensor]], int]:
     tensor_group: Dict[str, List[Tensor]] = {}
@@ -38,10 +35,6 @@ def load_h5(file_path: str) -> Tuple[Dict[str, List[Tensor]], int]:
                 for dset_name in grp.keys():
                     dset = grp[dset_name]
                     dsets.append(torch.from_numpy(dset[:]).share_memory_())
-                    total_samples += dset.attrs.get('numel', np.prod(dset.shape))
                 tensor_group[key] = dsets
 
-    num_keys = max(len(tensor_group), 1)
-    sample_per_key = total_samples // num_keys
-
-    return tensor_group, sample_per_key
+    return tensor_group