是否有带有 CUDA 统一 GPU-CPU 内存分支的 PyTorch？

Posted 2023-03-23

【中文标题】是否有带有 CUDA 统一 GPU-CPU 内存分支的 PyTorch？

【英文标题】：Is there a PyTorch with CUDA Unified GPU-CPU Memory fork?

【发布时间】：2020-10-29 03:28:43

【问题描述】：

因此，当一批一张图像需要 15GB 时，训练 DNN 模型可能会很痛苦。速度对我来说不是那么重要，但要适应更大的批次（和模型）。所以我想知道是否有带有 CUDA Unified Memory fork 的 PyTorch 或类似的东西来适应巨型模型（每个 GPU RAM 有 16gb，但在 CPU 端有 250gb 似乎相当合理）？

【问题讨论】：

似乎即使存在这样的解决方案，CPU 部分也会成为瓶颈，您可能不会从 GPU 中获得太多好处。您是否考虑过使用gradient accumulation？虽然如果你使用批量标准化层，这可能会变得很棘手。

【参考方案1】：

如果您不关心花费的时间，但需要大批量，您可以使用更慢的方法。假设您需要一批 128 个样本，但您的 gpu 内存仅适合 8 个样本。您可以创建 8 个样本的小批量，然后平均它们的梯度。

对于您评估的每小批 8 个样本，您将每个参数的 .grad 保留在 CPU 内存中。您为每个模型参数保留一个毕业生列表。在收集了 16 批 8 个样本（总共 128 个样本）的 grads 后，您可以对每个参数的梯度进行平均，并将结果放回每个参数的 .grad 属性中。

然后您可以调用优化器的.step()。这应该会产生与使用大量 128 个样本完全相同的结果。