| import os |
| import glob |
| import tarfile |
| import pickle |
| import zstandard |
| import torch |
| from torch.utils.data import IterableDataset, DataLoader |
|
|
| class DyMeshIterableDataset(IterableDataset): |
| def __init__(self, data_dir, rank=0, world_size=1): |
| super(DyMeshIterableDataset, self).__init__() |
| |
| self.tar_paths = sorted(glob.glob(os.path.join(data_dir, "shard_*.tar.zst"))) |
| |
| self.my_tar_paths = [p for i, p in enumerate(self.tar_paths) if i % world_size == rank] |
| print(f"[Rank {rank}] 负责处理 {len(self.my_tar_paths)} 个分片") |
|
|
| def __iter__(self): |
| worker_info = torch.utils.data.get_worker_info() |
| |
| if worker_info is None: |
| curr_paths = self.my_tar_paths |
| else: |
| per_worker = (len(self.my_tar_paths) + worker_info.num_workers - 1) // worker_info.num_workers |
| wid = worker_info.id |
| curr_paths = self.my_tar_paths[wid * per_worker : (wid + 1) * per_worker] |
|
|
| dctx = zstandard.ZstdDecompressor() |
| for tar_path in curr_paths: |
| try: |
| with open(tar_path, 'rb') as fh: |
| with dctx.stream_reader(fh) as reader: |
| with tarfile.open(fileobj=reader, mode='r|') as tar: |
| for member in tar: |
| if member.isfile(): |
| f = tar.extractfile(member) |
| if f: |
| data = pickle.load(f) |
| if isinstance(data, dict) and 'vertices' in data: |
| yield { |
| 'vertices': data['vertices'], |
| 'faces': data['faces'], |
| 'caption': data.get('caption', '无文本描述') |
| } |
| except Exception as e: |
| print(f"❌ 读取错误 {tar_path}: {e}") |
| continue |
|
|
| def collate_fn(batch): |
| """负责将数据组织成 Batch，格式与原始代码保持一致""" |
| return { |
| 'vertices': [torch.tensor(item['vertices'], dtype=torch.float32) for item in batch], |
| 'faces': [torch.tensor(item['faces'], dtype=torch.long) for item in batch], |
| 'caption': [item['caption'] for item in batch] |
| } |
|
|
| def get_mesh_dataloader(data_dir, batch_size=4, num_workers=2): |
| """统一的 DataLoader 获取接口""" |
| rank = int(os.environ.get("RANK", 0)) |
| world_size = int(os.environ.get("WORLD_SIZE", 1)) |
| |
| dataset = DyMeshIterableDataset(data_dir, rank=rank, world_size=world_size) |
| |
| return DataLoader( |
| dataset, |
| batch_size=batch_size, |
| collate_fn=collate_fn, |
| num_workers=num_workers, |
| pin_memory=True if torch.cuda.is_available() else False |
| ) |
