pytorch分布式系列3——分布式训练时，torch.utils.data.distributed.DistributedSampler做了什么？

最新推荐文章于 2025-03-18 09:40:37 发布

江洋大盗与鸭子

最新推荐文章于 2025-03-18 09:40:37 发布

阅读量1.5w

点赞数 56

分类专栏：深度学习文章标签： pytorch 神经网络分布式训练

本文链接：https://blog.csdn.net/searobbers_duck/article/details/115299691

版权

本文详细解释了torch.utils.data.distributed.DistributedSampler在PyTorch分布式训练中的作用，包括数据划分、epoch同步和batch大小调整。通过实例展示了不同配置下数据分布的变化以及如何解决数据平衡和批大小问题。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

分布式训练时，torch.utils.data.distributed.DistributedSampler做了什么？

试验用到的code

import os
import sys

import torch
import torch.nn as nn
import torch.distributed as dist
import torchvision

from torch.utils.data import Dataset, DataLoader

import numpy as np

class InnerDS(Dataset):
    def __init__(self, n=8):
        self.n = n

    def __len__(self):
        return self.n

    def __getitem__(self, item):
        np_img = np.random.rand(3,224,224)
        image = torch.from_numpy(np_img).float()
        label = np.random.randint(0,9)
        return image, label, item


local_rank = int(os.environ['LOCAL_RANK'])
world_size = int(os.environ['WORLD_SIZE'])
rank = int(os.environ['RANK'])

dist.init_process_group('nccl',world_size=world_size, rank=rank)


torch.cuda.set_device(local_rank)


# case 1
# ds = InnerDS(8)
# sampler = torch.utils.data.distributed.DistributedSampler(ds)
# dataloader = DataLoader(ds, batch_size=4, drop_last=True)

# case 2
# ds = InnerDS(8)
# sampler = torch.utils.data.distributed.DistributedSampler(ds)
# dataloader = DataLoader(ds, batch_size=4, sampler=sampler, drop_last=True)

# case 3
# ds = InnerDS(8)
# sampler = torch.utils.data.distributed.DistributedSampler(ds)
# dataloader = DataLoader(ds, batch_size=4, sampler=sampler, drop_last=True)

# case 4
# ds = InnerDS(6)
# sampler = torch.utils.data.distributed.DistributedSampler(ds)
# dataloader = DataLoader(ds, batch_size=4, sampler=sampler, drop_last=False)


# case 5
# ds = InnerDS(5)
# sampler = torch.utils.data.distributed.DistributedSampler(ds)
# dataloader = DataLoader(ds, batch_size=4, sampler=sampler, drop_last=False)

# case 6
# ds = InnerDS(10)
# sampler = torch.utils.data.distributed.DistributedSampler(ds)

最低0.47元/天解锁文章