PyTorch DDP流程和SyncBN、ShuffleBN

最新推荐文章于 2025-05-20 19:59:41 发布

原创

最新推荐文章于 2025-05-20 19:59:41 发布 · 413 阅读

4 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python

整体框架

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def train(local_rank, world_size):
    # 初始化进程组
    dist.init_process_group("nccl", rank=local_rank, world_size=world_size)
    torch.cuda.set_device(local_rank)
    
    # 数据加载器
    dataset = MyDataset()
    sampler = DistributedSampler(dataset)
    dataloader = DataLoader(dataset, sampler=sampler)
    
    # 模型与DDP包装
    model = MyModel().to(local_rank)
    model = DDP(model, device_ids=[local_rank])
    
    # 优化器
    optimizer = torch.optim.Adam(model.parameters())
    
    # 训练循环
    for epoch in range(epochs):
        sampler.set_epoch(epoch)
        for batch in dataloader:
            outputs = model(batch)
            loss = compute_loss(outputs)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
        
        # 保存模型（仅主进程）
        if local_rank == 0:
            torch.save(model.module.state_dict(), "model.pth")
    
    # 清理进程组
    dist.destroy_process_group()