小编ymf*_*mfj的帖子

我想知道如何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, …