存储由torch DataLoader加载的CUDA或CPU张量哪个更好？

Question

我正在开展一个项目，目标是在

多个 GPU

上训练 PyTorch 模型。

我的输入数据存储在每个训练示例的单独文件中，在预处理过程中，我使用

torch.save

方法将它们保存到

.pt

文件中。稍后，我使用

DataLoader

加载这些文件，我想在其中设置

num_workers > 0

来加速 进程。然而，似乎只有当输入数据在CPU上时才能将num_workers

设置为

>0

。

我的问题是：我应该保存

CUDA tensors

 并只使用 

num_workers=0

，还是应该存储 

CPU tensors

，设置 

num_workers > 0

，然后将批次作为一个整体移动到 GPU？

我不确定哪种方法在多个 GPU 上的训练速度（时间）上更有效。任何有关此事的见解或最佳实践将不胜感激。

Answer 1

num_workers>0

 允许在 CPU 上并行化，而不是在 GPU 上并行化。一般来说，最佳实践是在 CPU 上跨多个进程（使用 

num_workers

）执行数据加载和预处理。

通过

DataLoader

 整理批次后，将整个批次传输到训练循环中的 GPU。

要在多个 GPU 上进行训练，您可以使用 PyTorch 的

分布式数据并行 (DDP) 或数据并行 (DP) 包装器在 GPU 上分布模型和数据。