File size: 3,032 Bytes
f0fc238 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | 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")))
# 根据 Rank 进行数据切分
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()
# 根据 Worker 进行数据切分
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
) |