如何使用 PyTorch 的 DataLoader 确保批次包含来自所有工作人员的样本？

Question

我想知道如何torch.utils.data.DataLoader在 PyTorch 中使用，尤其是在多工人情况下。

我发现一个批次的输出DataLoader总是来自一个工人。我预计 DataLoader 中有一个队列，用于存储来自所有工作人员的数据，而 DataLoader 将它们在队列中打乱以输出随机批处理数据。我认为这是tf.data.DatasetTensorflow 中的方法。我们可以在 PyTorch 中实现类似的功能吗？我想Tfrecord通过使用多个工作人员从大序列化文件（如）加载数据集。在这种情况下，在一批中混合源文件，这意味着混合工作器的源，很重要。

请参考以下代码：

import random
import time

import torch


class MyDataset(torch.utils.data.Dataset):
    def __len__(self):
        return 50

    def __getitem__(self, idx):
        info = torch.utils.data.get_worker_info()

        time.sleep(random.uniform(0, 1))
        print("[{}]:{}".format(info.id, idx))
        return idx, info.id


if __name__ == '__main__':
    dataset = MyDataset()
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=5, shuffle=False, num_workers=2)
    for batch in dataloader:
        print(batch)

输出：

[0]:0
[1]:5
[0]:1
[1]:6
[0]:2
[0]:3
[1]:7
[0]:4
[tensor([0, 1, 2, 3, 4]), tensor([0, 0, 0, 0, 0])]
[1]:8
[1]:9
[tensor([5, 6, 7, 8, 9]), tensor([1, 1, 1, 1, 1])]
[0]:10
[0]:11
[1]:15
[1]:16
[0]:12
[1]:17
...

在这里，[0, 1, 2, 3, 4]并[0, 0, 0, 0, 0]在[tensor([0, 1, 2, 3, 4]), tensor([0, 0, 0, 0, 0])]平均，此批包括指数0个到4个数据来自工人的ID 0。请注意，这shuffle=True并不能解决仅更改数据索引的问题。

在这种情况下，我希望得到一个批次，如：[tensor([0, 5, 1, 6, 2]), tensor([0, 1, 0, 1, 0])]。

Answer 1

请注意，指定了batch_size的多workerDataLoader将并行加载多个批次，因此本质上一个批次始终来自一个worker。但是，我通过执行以下操作已经实现了接近您要求的目标：

1. 将batch size设置为1，这样每个工人一次只能产出一个样本

2. 编写一个后台进程，循环访问 DataLoader，一次获取 1 个样本并将其插入队列。这使得可以在队列中按不同的顺序排列样本，而不是使用特定于工作人员的批次

3. 有一个批处理机制，例如collate_fn从队列中获取等于批量大小的样本并将其提供给模型

如果您想在批量创建中更加具体，例如从特定工作人员中挑选特定样本，您可以有多个队列。应修改您的整理程序以考虑多个队列并从中进行选择。但我怀疑是否需要这种特殊性。