ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor(data): 修改文件加载方案

Browse Source
This commit is contained in:
ViperEkura
2026-02-22 21:14:10 +08:00
parent 0ca4871e80
commit 582d4ae9a7
4 changed files with 67 additions and 246 deletions
Download Patch File Download Diff File Expand all files Collapse all files
khaosz/data/dataset.py
+10 -10
View File
@@ -1,18 +1,17 @@
import h5py
import torch
import bisect
from abc import ABC, abstractmethod
from torch import Tensor
from torch.utils.data import Dataset
from khaosz.data.mmap import MmapFileHandler
from khaosz.data.file import load_h5
from typing import Callable, List, Dict, Literal, Optional, Union
Seg = List[Tensor]
MultiSeg = Dict[str, Seg]
class BaseSegmentFetcher:
def __init__(self, segments: Seg):
def __init__(self, segments: List[Tensor]):
self.segments = segments
self.cum_lengths = []
total = 0
@@ -37,20 +36,21 @@ class BaseSegmentFetcher:
prev_cum = self.cum_lengths[i - 1] if i > 0 else 0
start = max(begin_idx - prev_cum, 0)
end = min(end_idx - prev_cum, len(self.segments[i]))
result_segments.append(self.segments[i][start:end])
data = self.segments[i][start:end]
result_segments.append(data)
return torch.cat(result_segments, dim=0)
class MultiSegmentFetcher:
def __init__(self, muti_segments: MultiSeg):
def __init__(self, muti_segments: Dict):
self.muti_keys = list(muti_segments.keys())
self.muti_fetchers = {
key: BaseSegmentFetcher(segments)
for key, segments in muti_segments.items()
}
def key_fetch(self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]) -> Union[Tensor, Seg]:
def key_fetch(self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]) -> Dict:
fetch_dict = {}
keys = [keys] if isinstance(keys, str) else keys
@@ -61,20 +61,20 @@ class MultiSegmentFetcher:
return fetch_dict if len(keys) > 1 else fetch_dict[keys[0]]
def fetch_data(self, begin_idx: int, end_idx: int) -> Union[Tensor, Seg]:
def fetch_data(self, begin_idx: int, end_idx: int) -> Dict:
return self.key_fetch(begin_idx, end_idx, self.muti_keys)
class BaseDataset(Dataset, ABC):
def __init__(self, window_size: int, stride: int):
super().__init__()
self.segments: MultiSeg = {}
self.segments = {}
self.window_size = window_size
self.stride = stride
self.total_samples = None
def load(self, load_path: str):
self.segments, self.total_samples = MmapFileHandler.load(load_path)
self.segments, self.total_samples = load_h5(load_path)
self.fetcher = MultiSegmentFetcher(self.segments)
def get_index(self, index: int) -> int:
khaosz/data/file.py
+44
View File
@@ -0,0 +1,44 @@
import os
import h5py
import numpy as np
import torch
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:
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,
compression='gzip',
compression_opts=4,
shuffle=True
)
dset.attrs['numel'] = tensor.numel()
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
num_keys = max(len(tensor_group), 1)
sample_per_key = total_samples // num_keys
return tensor_group, sample_per_key
khaosz/data/mmap.py
-82
View File
@@ -1,82 +0,0 @@
import os
import json
import torch
from torch import Tensor
from typing import List, Dict, Tuple
class MmapFileHandler:
"""
json metadata like this:
```
[
{"file_name": "file1.pt", "size": 1000, "key": "key1"},
{"file_name": "file2.pt", "size": 2000, "key": "key2"}
...
]
```
files like:
```
folder_path:
- metadata.json
- file1.pt
- file2.pt
...
```
"""
META_DATA = "metadata.json"
@staticmethod
def load(root_path: str, shared: bool=True) -> Tuple[Dict[str, List[Tensor]], int]:
metadata_list = []
tensor_group: Dict[str, List[Tensor]] = {}
file_mapper_path = os.path.join(root_path, MmapFileHandler.META_DATA)
if not os.path.exists(file_mapper_path):
raise FileNotFoundError(f"File mapper not found: {file_mapper_path}")
with open(file_mapper_path, "r") as f:
metadata_list = json.load(f)
for metadata in metadata_list:
file_key = metadata["key"]
file_name = metadata["file_name"]
file_path = os.path.join(root_path, file_name)
elm = torch.load(file_path, map_location="cpu", mmap=shared)
if file_key not in tensor_group:
tensor_group[file_key] = []
tensor_group[file_key].append(elm)
num_samples = sum(metadata["size"] for metadata in metadata_list)
num_keys = max(len(set(metadata['key'] for metadata in metadata_list)), 1)
sample_per_key = num_samples // num_keys
return tensor_group, sample_per_key
@staticmethod
def save(save_path: str, mmap_shared_group: Dict[str, List[Tensor]]) -> None:
os.makedirs(save_path, exist_ok=True)
metadata_list = []
for segment_key, segment_tensors in mmap_shared_group.items():
for idx, tensor in enumerate(segment_tensors):
try:
with open(os.path.join(save_path, f"{segment_key}_{idx}.pt"), "wb") as f:
torch.save(tensor.contiguous().cpu(), f)
except Exception as e:
raise RuntimeError(f"Error saving tensor: {e}")
metadata_list.append({
"file_name": f"{segment_key}_{idx}.pt",
"size": tensor.numel(),
"key": segment_key
})
metadata_path = os.path.join(save_path, MmapFileHandler.META_DATA)
with open(metadata_path, "w") as f:
json.dump(metadata_list, f)