fix: 修复一些运行时问题

khaosz/data/checkpoint.py

@@ -12,14 +12,12 @@ from khaosz.parallel.setup import get_rank
 class Checkpoint:
     def __init__(
         self,
-        optimizer_state_dict: Dict[str, Any],
-        scheduler_state_dict: Optional[Dict[str, Any]] = None,
+        state_dict: Dict[str, Any],
         epoch: int = 0,
         iteration: int = 0,
         metrics: Optional[Dict[str, list]] = None,
     ):
-        self.optimizer_state_dict = optimizer_state_dict
-        self.scheduler_state_dict = scheduler_state_dict
+        self.state_dict = state_dict
         self.epoch = epoch
         self.iteration = iteration
         self.metrics = metrics or {}
@@ -46,12 +44,8 @@ class Checkpoint:
             if save_metric_plot and self.metrics:
                 self._plot_metrics(str(save_path))
 
-        state_dict = {
-            "optimizer": self.optimizer_state_dict,
-            "scheduler": self.scheduler_state_dict
-        }
-        with open(save_path / f"state_dict_rank_{get_rank()}.pt", "wb") as f:
-            torch.save(state_dict, f)
+            with open(save_path / f"state_dict.pt", "wb") as f:
+                torch.save(self.state_dict, f)
 
     @classmethod
     def load(
@@ -72,7 +66,7 @@ class Checkpoint:
             dist.broadcast_object_list(meta_list, src=0)
             meta = meta_list[0]
 
-        with open(save_path / f"state_dict_rank_{get_rank()}.pt", "rb") as f:
+        with open(save_path / f"state_dict.pt", "rb") as f:
             state_dict = torch.load(f)
 
         return cls(

khaosz/data/dataset.py

@@ -1,4 +1,3 @@
-import h5py
 import torch
 import bisect
 
@@ -78,8 +77,10 @@ class BaseDataset(Dataset, ABC):
         self.fetcher = MultiSegmentFetcher(self.segments)
         
     def get_index(self, index: int) -> int:
-        begin_idx = min(index * self.stride, self.total_samples - self.window_size - 1)
-        end_idx = begin_idx + self.window_size
+        assert self.total_samples > self.window_size
+        
+        begin_idx = min(index * self.stride, self.total_samples - 1 - self.window_size)
+        end_idx = min(begin_idx + self.window_size, self.total_samples - 1)
         
         return begin_idx, end_idx
         
@@ -91,7 +92,7 @@ class BaseDataset(Dataset, ABC):
         assert self.total_samples is not None
         if self.total_samples <= self.window_size:
             return 0
-        return self.total_samples // self.stride + 1
+        return (self.total_samples - 1 - self.window_size) // self.stride + 1
     
 
 class SeqDataset(BaseDataset):

khaosz/data/file.py

@@ -2,6 +2,8 @@ import os
 import h5py
 import numpy as np
 import torch
+
+from pathlib import Path
 from torch import Tensor
 from typing import Dict, List, Tuple
 
@@ -17,10 +19,7 @@ def save_h5(file_path: str, tensor_group: Dict[str, List[Tensor]]):
                 arr = tensor.cpu().numpy()
                 dset = grp.create_dataset(
                     f'data_{idx}',
-                    data=arr,
-                    compression='gzip',
-                    compression_opts=4,
-                    shuffle=True
+                    data=arr
                 )
                 dset.attrs['numel'] = tensor.numel()
 
@@ -28,15 +27,19 @@ def load_h5(file_path: str) -> Tuple[Dict[str, List[Tensor]], int]:
     tensor_group: Dict[str, List[Tensor]] = {}
     total_samples = 0
 
-    with h5py.File(file_path, 'r') as f:
-        for key in f.keys():
-            grp = f[key]
-            dsets = []
-            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
+    root_path = Path(file_path)
+    h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
+    
+    for h5_file in h5_files:
+        with h5py.File(h5_file, 'r') as f:
+            for key in f.keys():
+                grp = f[key]
+                dsets = []
+                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