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

Question

1. 如果num_workers是2，那么它会将2批次放入RAM并发送其中1个到GPU，还是会将3批次放入RAM然后发送其中1个到GPU？
2. 当工作进程数大于CPU核心数时，实际上会发生什么？我尝试过，一切正常，但它是如何工作的？（我认为我可以选择的最大工作进程数就是核心数）。
3. 如果我将num_workers设置为3，并且在训练期间GPU内存中没有批次数据，则主进程会等待其工作进程读取批次数据，还是会读取单个批次（而不等待工作进程）？

- floyd

---

可能会感兴趣的内容：https://discuss.pytorch.org/t/guidelines-for-assigning-num-workers-to-dataloader/813 - Charlie Parker

Answer 1

1. 当 num_workers>0 时，只有这些工作进程将检索数据，主进程不会。 所以当 num_workers=2 时，最多有2个工作进程同时将数据放入RAM中，而不是3个。
2. 我们的CPU通常可以运行大约100个进程，而这些工作进程并没有任何特殊之处，因此拥有比CPU核心更多的工作进程是可以的。但它是否高效取决于CPU核心在执行其他任务时的繁忙程度、CPU的速度、硬盘的速度等等。简而言之，情况很复杂，所以将工作进程设置为核心数是一个好的经验法则，仅此而已。
3. 不是的。请记住，DataLoader 不只是从当前可用的RAM中随机返回，它使用 batch_sampler来决定下一批要返回的内容。每个批次都分配给一个工作进程，主进程将等待分配的工作进程检索到所需的批次。

最后澄清一下，DataLoader 的工作不是直接发送任何内容到GPU，您必须显式调用 cuda() 进行操作。

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

- Shihab Shahriar Khan

---

仅对最后一句话进行备注 - 在Dataset对象中调用.cuda()可能不是一个好主意，因为它将不得不单独将每个样本（而不是批次）移动到GPU上，从而产生大量开销。 - psarka

我还想补充一点，在Windows上将工作线程数设置为大于0可能会导致错误（参见https://discuss.pytorch.org/t/errors-when-using-num-workers-0-in-dataloader/97564/3）。 - Marine Galantin

我没有测试过这个，但你可能可以在collate_fn函数中将数据移动到gpu上。假设这个函数也是并行的，那么它可以加快速度。潜在的问题是现在你在gpu上有大于等于n_workers批次的数据，所以内存可能会受到限制。 - mkohler