PyTorch 数据加载器中的“工人数量”参数实际上是如何工作的？

Posted 2023-02-16

【中文标题】PyTorch 数据加载器中的“工人数量”参数实际上是如何工作的？

【英文标题】：How does the "number of workers" parameter in PyTorch dataloader actually work?

【发布时间】：2019-05-28 15:00:46

【问题描述】：

如果num_workers 为 2，这是否意味着它会将 2 个批次放入 RAM 中并将其中 1 个批次发送到 GPU，还是将 3 个批次放入 RAM 中然后将其中 1 个批次发送到 GPU？李> 当工作人员的数量高于 CPU 内核的数量时，实际会发生什么？我试过了，它工作得很好，但它是如何工作的？ （我认为我可以选择的最大工人数是核心数）。 如果我将num_workers 设置为 3，并且在训练期间内存中没有用于 GPU 的批次，主进程是等待其工作人员读取批次还是读取单个批次（不等待工人）？

【问题讨论】：

可能感兴趣：discuss.pytorch.org/t/…

【参考方案1】：

num_workers>0 时，只有这些工作人员会检索数据，主进程不会。因此，当num_workers=2 时，您最多有 2 个工作人员同时将数据放入 RAM，而不是 3 个。 好吧，我们的 CPU 通常可以毫无问题地运行 100 个进程，而且这些工作进程并不特殊，因此拥有比 CPU 内核更多的工作进程是可以的。但它有效率吗？这取决于您的 cpu 核心对其他任务的繁忙程度、cpu 速度、硬盘速度等。简而言之，它很复杂，因此将工作人员设置为核心数量是一个很好的经验法则，仅此而已。 不。请记住，DataLoader 不只是从 RAM 中的可用内存中随机返回，它使用batch_sampler 来决定下一个要返回的批次。每个批次都分配给一个工人，主进程将等待，直到分配的工人检索到所需的批次。

最后澄清一下，将任何内容直接发送到 GPU 不是 DataLoader 的工作，您为此明确调用 cuda()。

编辑：不要在Dataset 的__getitem__() 方法中调用cuda()，请查看@psarka 的评论以了解原因

【讨论】：

只是对最后一句的注释 - 在 Dataset 对象中调用 .cuda() 可能不是一个好主意，因为它必须将每个样本（而不是批次）移动到单独使用 GPU，会产生大量开销。

我还想补充一点，在 windows 上设置高于 0 的工作人员数量可能会导致错误（参见discuss.pytorch.org/t/…）。