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

Question

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

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

我想为我的课程研究对象检测算法。许多深度学习架构需要大容量的 GPU 内存，所以我的机器无法训练这些模型。我尝试通过将每一层加载到 GPU 然后将其加载回来来处理图像：

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

但它似乎不是很有效。我想知道在使用很少的 GPU 内存的同时训练大型深度学习模型是否有任何提示和技巧。提前致谢！

编辑：我是深度学习的初学者。如果这是一个愚蠢的问题，请道歉:)

Answer 1

虽然，

    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 内存不足的原因并重新启动内核以避免错误再次发生（就像我在我的情况下所做的那样）。

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

Answer 2

以迭代方式将批次发送到 CUDA，并制作小批量。不要一开始就将所有数据一次发送到 CUDA。相反，请按以下方式操作：

for e in range(epochs):
    for images, labels in train_loader:   
        if torch.cuda.is_available():
            images, labels = images.cuda(), labels.cuda()   
        # blablabla  

您也可以dtypes使用更少的内存。例如，torch.float16或torch.half。

Answer 3

只需减少批量大小，它就会起作用。当我在训练时，它给出了以下错误：

CUDA 内存不足。尝试分配 20.00 MiB（GPU 0；10.76 GiB 总容量；4.29 GiB 已分配；10.12 MiB 空闲；PyTorch 总共保留 4.46 GiB）

我使用的是 32 的批量大小。所以我只是将它更改为 15，它对我有用。

Answer 4

大部分内容都涵盖了，还会补充一点。

如果 torch 给出错误“尝试分配 2 MiB”等，则这是一条误导性消息。实际上，CUDA 耗尽了训练模型所需的总内存。您可以减小批量大小。比如说，即使批量大小为 1 不起作用（当您使用大量序列训练 NLP 模型时会发生这种情况），请尝试传递较少的数据，这将帮助您确认您的 GPU 没有足够的内存来训练模型。

此外，如果您想重新训练模型，则必须再次完成垃圾收集和清理缓存部分。

Answer 5

尽量不要把你的毕业生拖得太远。

当我试图总结所有批次的损失时，我遇到了同样的错误。

loss =  self.criterion(pred, label)

total_loss += loss

然后我使用 loss.item 而不是需要 grads 的 loss，然后解决了问题

loss =  self.criterion(pred, label)

total_loss += loss.item()

下面的解决方案归功于yuval reina in the kaggle question

此错误与 GPU 内存有关，与一般内存无关 => @cjinny 注释可能不起作用。
你使用 TensorFlow/Keras 还是 Pytorch？
尝试使用较小的批量大小。
如果您使用 Keras，请尝试减小一些隐藏层的大小。
如果您使用 Pytorch：
您是否一直将所有训练数据保存在 GPU 上？
确保不要将渐变拖得太远
检查隐藏层的大小

Answer 6

如果您完成训练并且只想使用图像进行测试，请确保在开头添加 with torch.no_grad() 和 m.eval() ：

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

这看起来似乎很明显，但它对我的情况有效。我试图使用 BERT 将句子转换为嵌入表示。由于 BERT 是一个预先训练的模型，我不需要保存所有梯度，而且它们会消耗所有 GPU 内存。