如何解决“RuntimeError: CUDA Out of memory”问题

作者：Nitin Kishore

编译：ronghuaiyang

导读

一些可以尝试的解决“RuntimeError: CUDA Out of memory”的方案。

当遇到这个问题时，你可以尝试一下这些建议，按代码更改的顺序递增：

减少“batch_size”
降低精度
按照错误说的做
清除缓存
修改模型/训练

在这些选项中，如果你使用的是预训练模型，则最容易和最有可能解决问题的选项是第一个。

修改batchsize

如果你是在运行现成的代码或模型，则最好的做法是减小batchsize。减半，然后继续减半，直到没有错误为止。

但是，如果在此过程中，你发现自己将batchsize大小设置为 1 并且仍然无济于事，那么就还有其他问题，如果可以修复它，那么模型训练可以在更大的batchsize下工作。

降低精度

如果你用的是 Pytorch-Lightning，你也可以尝试将精度更改为“float16”。这可能会带来诸如预期的 Double 和 Float 张量之间的不匹配等问题，但它可以节省很多内存的，并且在性能上有一个非常轻微的权衡，使其成为一个可行的选择。

这第三种选择 ——

按照错误信息去做

可以使用以下命令完成此操作。如果你使用的是 Windows 计算机，则可以使用 set 而不是 export

export PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.6,max_split_size_mb:128

如果你在Jupyter或Colab笔记本上，在发现RuntimeError: CUDA out of memory后。你需要重新启动kernel。

使用多 GPU 系统时，我建议使用CUDA_VISIBLE_DEVICES 环境变量来选择要使用的 GPU。

$ export CUDA_VISIBLE_DEVICES=0  (OR)
$ export CUDA_VISIBLE_DEVICES=1  (OR) 
$ export CUDA_VISIBLE_DEVICES=2,4,6  (OR)# This will make the cuda visible with 0-indexing so you get cuda:0 even if you run the second one.

或者，在python代码中设置：

import os
os.environ['CUDA_VISIBLE_DEVICES']='2, 3'

一些stack overflow的帖子会让你尝试将这些行添加到你的代码中 ——

释放Cache

要弄清楚你的模型在 cuda 上占用了多少内存，你可以尝试：

import gc
def report_gpu():
   print(torch.cuda.list_gpu_processes())
   gc.collect()
   torch.cuda.empty_cache()

如果你调用python的垃圾收集，并调用pytorch的清空缓存，这基本上应该让你的GPU恢复到一个干净的状态，不使用超过它需要的内存，当你开始训练下一个模型时，不必重新启动kernel。

import gc
gc.collect()
torch.cuda.empty_cache()

虽然torch.cuda.empty_cache() 或gc.collect() 可以释放CUDA内存，但显然不能释放的内存返回到Python中。因此，不要把希望寄托在这些脚本上。对于JupyterLab或Colab来说，这种方式是有效的。下面是如何使用这些代码的例子：

我们马上去看一下 .detach()和.cpu()：

这通常仅适用于notebooks 和 ipython

让我们看看这些之外的替代方法

使用koila python包

以下是如何使用koila的例子

pip install koila
# Wrap the input tensor and label tensor.
# If a batch argument is provided, that dimension of the tensor would be treated as the batch.
# In this case, the first dimension (dim=0) is used as batch's dimension.

(input, label) = lazy(input, label, batch=0)

如果这些都没有帮助，解决问题的唯一方法是找出：

是什么在使用内存?

与流行的看法相反，你不需要更大的GPU来训练更大的模型，你可以简单地使用梯度累积。我们稍后会讨论这个问题。我们来剖析一下错误消息，因为这通常是一个很好的提示。

这条信息告诉你总共有15.78G的GPU内存。因此，你可以检查一下的数据和模型有多大，因为你需要将它们移动到 GPU中。如果它超过了总容量，则无法在该计算机上运行，你需要将数据分块，并在 CPU 和 GPU 之间持续移动。

减小图像尺寸也有帮助，如果它说你不能使用 x MiB，因为你只有一点内存可用，找出其他进程也在使用 GPU 并释放该空间。通过运行以下命令查找 python 进程的 PID：

nvidia-smi

杀掉

sudo kill -9 pid

修改模型/训练循环

现在我们进入了最后阶段。一切都试过了，没有任何东西在使用 GPU 内存，你编写的代码很可能将很多东西都推送到 GPU 上了。你想要它用于大型矩阵乘法，但对于像指标计算和日志记录这样简单的东西，你可以在CPU上做这些。所以把这些都从GPU上去掉。

Loss, Preds, Targets

不要保存整个tensor，当你在epoch结束需要汇总损失的时候，使用loss.item()。

使用preds.detach().cpu()从 GPU 中删除预测和目标。这些都是很重的东西。如果你只保留它们用于日志记录，则无需将它们保留在内存中。

在这里要小心。如果你发现你的损失在各个epoch中是恒定的，那可能是因为你分离了计算图的一部分，而反向传播没有办法返回更新值。因此，请弄清楚在代码中的哪个点可以执行上述 2 个步骤。

我们再说一下另外一个技巧：

梯度累积

减小batchsize大小是避免内存问题的一种方法，但是，batchsize越小，batch与batch之间的波动性就越大。因此，训练的动态会有所不同。你不希望继续为不同batchsize大小的结构查找一组不同的超参数。

你可以使用另一个称为accum的参数来“累积梯度”，方法是定义累积的梯度batch数。由于我们在accum批次上添加这些梯度，因此我们将batch_size除以相同的数字。

batch_size//accum

累积可以通过使用回调GradientAccumulation来完成

batch_size = 64
accum=2
# Data loader , change the batchsize parmeter to  bs = 64//accum
cbs = GradientAccumulation(64) if accum else []

我们可以做的是找到一种方法来运行 32 条数据，但让它一次表现得像 64 条数据。在一般的训练循环中，在执行loss.backward() 之前，你需要将梯度归零，如果不将梯度归零，梯度将进行累积。

因此，如果你在不将梯度归零的情况下进行 2 个半批次的训练，它们的梯度回累积，最终得到以目标有效批次大小相同的梯度。在训练循环中，我们需要使用计数器根据小batchsize大小进行更新，一旦它达到预设的目标，那就是我们将梯度归零时。在那之前，他们只是通过loss.backward()而不断积累。

英文原文：https://medium.com/@snk.nitin/how-to-solve-cuda-out-of-memory-error-850bb247cfb2