File size: 3,836 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | import os
import glob
import tarfile
import pickle
import zstandard
import torch
import torch.distributed as dist
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)
yield {'vertices': data['vertices'], 'faces': data['faces'], 'caption': data.get('caption', '无')}
except Exception: continue
def get_mesh_dataloader(data_dir, batch_size=4, num_workers=2):
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=lambda x: x, num_workers=num_workers, pin_memory=True)
def main():
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
global_rank = dist.get_rank()
world_size = dist.get_world_size()
torch.cuda.set_device(local_rank)
dataloader = get_mesh_dataloader("/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/data/shards/", batch_size=40, num_workers=8)
local_batch_count = 0
local_sample_count = 0
for batch_idx, batch in enumerate(dataloader):
local_batch_count += 1
local_sample_count += len(batch) # 统计实际样本数
if batch_idx % 50 == 0:
print(f"[Rank {global_rank}] 进度: 处理至第 {local_batch_count} 个 batch (累计样本: {local_sample_count})")
dist.barrier()
# 汇总 Batch 数量
batch_counts = torch.zeros(world_size, dtype=torch.long, device=f"cuda:{local_rank}")
batch_counts[global_rank] = local_batch_count
dist.all_reduce(batch_counts, op=dist.ReduceOp.SUM)
# 汇总 Sample 数量
sample_counts = torch.zeros(world_size, dtype=torch.long, device=f"cuda:{local_rank}")
sample_counts[global_rank] = local_sample_count
dist.all_reduce(sample_counts, op=dist.ReduceOp.SUM)
if global_rank == 0:
print("\n" + "="*50)
print("📊 分布式训练统计报告")
for i in range(world_size):
print(f" Rank {i}: 处理了 {batch_counts[i].item()} 个 Batch, 共 {sample_counts[i].item()} 个样本")
print("-" * 50)
print(f" 全局总 Batch 数: {batch_counts.sum().item()}")
print(f" 全局总样本数 : {sample_counts.sum().item()}")
print("="*50 + "\n")
dist.destroy_process_group()
if __name__ == "__main__":
main() |