PyTorch Lightning中使用DDP时如何保留数据集顺序？

要在PyTorch Lightning多GPU训练中保留数据输入模型的顺序，你可以通过自定义数据采样器覆写默认的数据集拆分逻辑，以下是具体实现思路和代码示例：

核心思路 #

默认分布式采样器会将数据集按GPU数量切分为连续分段，要保留全局顺序，需让每个GPU按固定间隔采样数据而非取连续块。比如2个GPU时，GPU0取索引0、2、4...，GPU1取索引1、3、5...，这样全局批次的顺序就能和原始数据集一致。

具体实现 #

1. 自定义分布式采样器 #

继承torch.utils.data.DistributedSampler，重写__iter__方法生成按间隔采样的索引：

import torch
from torch.utils.data import DistributedSampler

class OrderedDistributedSampler(DistributedSampler):
    def __iter__(self):
        # 获取全局数据集索引
        indices = list(range(len(self.dataset)))
        # 按GPU数量和当前进程ID生成间隔采样的索引
        indices = indices[self.rank::self.num_replicas]
        return iter(indices)

2. 在LightningDataModule中使用自定义采样器 #

在train_dataloader等方法中指定自定义采样器：

from pytorch_lightning import LightningDataModule

class CustomDataModule(LightningDataModule):
    def __init__(self, dataset, batch_size=32):
        super().__init__()
        self.dataset = dataset
        self.batch_size = batch_size

    def train_dataloader(self):
        sampler = OrderedDistributedSampler(self.dataset)
        return torch.utils.data.DataLoader(
            self.dataset,
            batch_size=self.batch_size,
            sampler=sampler,
            num_workers=4
        )

3. 注意事项 #

• 确保训练时使用分布式模式启动（如torchrun或pl.Trainer(strategy="ddp")）
• 验证全局顺序时，可在训练日志中打印每个batch的样本索引，确认是否和原始数据集顺序一致
• 自定义采样器会保持每个GPU的样本数量尽可能均衡，避免负载不均

内容的提问来源于stack exchange，提问作者malfonsoarquimea