如何避免 PyTorch 中的“CUDA 内存不足”

我认为对于 GPU 内存较低的 PyTorch 用户来说，这是一个非常常见的消息：

RuntimeError: CUDA out of memory. Tried to allocate X MiB (GPU X; X GiB total capacity; X GiB already allocated; X MiB free; X cached)

我尝试通过将每一层加载到 GPU 然后加载回来来处理图像：

for m in self.children():
    m.cuda()
    x = m(x)
    m.cpu()
    torch.cuda.empty_cache()

但似乎效果不是很大。我想知道是否有任何提示和技巧可以在使用少量 GPU 内存的情况下训练大型深度学习模型。

Although

import torch
torch.cuda.empty_cache()

为清除占用的cuda内存提供了一个很好的选择，我们还可以通过使用手动清除未使用的变量，

import gc
del variables
gc.collect()

但在使用这些命令后，错误可能会再次出现，因为 pytorch 实际上并没有清除内存，而是清除了对变量占用的内存的引用。 因此，重新启动内核后减小batch_size并找到最佳batch_size是最好的选择（但有时不是一个非常可行的选择）。

深入了解 GPU 内存分配的另一种方法是使用：

torch.cuda.memory_summary(device=None, abbreviated=False)

其中，两个参数都是可选的。这提供了内存分配的可读摘要，并允许您找出 CUDA 内存不足的原因并重新启动内核以避免再次发生错误（就像我在我的例子中所做的那样）。

迭代传递数据可能会有所帮助，但改变网络层的大小或分解它们也被证明是有效的（因为有时模型也会占用大量内存，例如在进行迁移学习时）。