如何在PyTorch中构建可复用且带洗牌的并联多DataLoader

要实现可跨epoch复用且自带同步洗牌逻辑的DataLoader，核心是基于PyTorch的Dataset类自定义组合数据集，而非直接拼接DataLoader迭代器。以下是具体实现步骤：

1. 自定义组合数据集类 #

创建继承自torch.utils.data.Dataset的类，将三个原始数据集打包，确保同一索引对应的样本被同时取出，保证样本配对的一致性：

import torch
from torch.utils.data import Dataset, DataLoader

class CombinedDataset(Dataset):
    def __init__(self, dataset_a, dataset_b, dataset_c):
        self.dataset_a = dataset_a
        self.dataset_b = dataset_b
        self.dataset_c = dataset_c
        # 校验三个数据集长度一致
        assert len(self.dataset_a) == len(self.dataset_b) == len(self.dataset_c), "三个数据集长度必须相同"
    
    def __len__(self):
        return len(self.dataset_a)
    
    def __getitem__(self, idx):
        # 根据索引取出对应位置的三个样本，打包为字典
        return {
            'a': self.dataset_a[idx],
            'b': self.dataset_b[idx],
            'c': self.dataset_c[idx]
        }

2. 创建可复用的DataLoader D #

基于自定义组合数据集生成DataLoader，利用PyTorch原生的洗牌机制实现epoch间的自动乱序：

# 从现有DataLoader中提取原始数据集（若你已有原始Dataset可直接使用）
dataset_a = A.dataset
dataset_b = B.dataset
dataset_c = C.dataset

# 初始化组合数据集
combined_dataset = CombinedDataset(dataset_a, dataset_b, dataset_c)

# 创建最终的DataLoader D，沿用原DataLoader的批量大小、工作进程数等参数
D = DataLoader(
    combined_dataset,
    batch_size=A.batch_size,
    shuffle=True,  # 开启自动洗牌，每个epoch都会重新打乱样本顺序
    num_workers=A.num_workers,
    collate_fn=A.collate_fn  # 若原DataLoader有自定义批量拼接逻辑，可直接复用
)

3. 使用方式 #

现在D可以跨epoch重复迭代，每次迭代返回符合要求的字典批量：

num_epochs = 10
for epoch in range(num_epochs):
    for batch_dict in D:
        # 取出各数据集的批量样本
        batch_a = batch_dict['a']
        batch_b = batch_dict['b']
        batch_c = batch_dict['c']
        # 后续模型训练/推理逻辑

方案优势 #

• 跨epoch复用：DataLoader每次迭代时会重新从Dataset生成迭代器，无需手动重新初始化，直接重复循环D即可完成多epoch训练。
• 同步洗牌：通过同一索引取数，确保三个数据集的样本始终配对，且PyTorch自动处理epoch间的洗牌逻辑，无需手动统一随机种子。
• 兼容性强：可完全复用原DataLoader的批量大小、多进程、自定义拼接函数等参数，与原有训练流程无缝衔接。

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