一些可以尝试的解决“RuntimeError: CUDA Out of memory”的方案。
当遇到这个问题时,你可以尝试一下这些建议,按代码更改的顺序递增:
- 减少“batch_size”
- 降低精度
- 按照错误说的做
- 清除缓存
- 修改模型/训练
在这些选项中,如果你使用的是预训练模型,则最容易和最有可能解决问题的选项是第一个。
修改batchsize
如果你是在运行现成的代码或模型,则最好的做法是减小batchsize。减半,然后继续减半,直到没有错误为止。
但是,如果在此过程中,你发现自己将batchsize大小设置为 1 并且仍然无济于事,那么就还有其他问题,如果可以修复它,那么模型训练可以在更大的batchsize下工作。
降低精度
如果你用的是 Pytorch-Lightning,你也可以尝试将精度更改为“float16”。这可能会带来诸如预期的 Double 和 Float 张量之间的不匹配等问题,但它可以节省很多内存的,并且在性能上有一个非常轻微的权衡,使其成为一个可行的选择。
这第三种选择 ——
按照错误信息去做
可以使用以下命令完成此操作。如果你使用的是 Windows 计算机,则可以使用 set 而不是 export
如果你在Jupyter或Colab笔记本上,在发现
RuntimeError: CUDA out of memory
后。你需要重新启动kernel。
使用多 GPU 系统时,我建议使用CUDA_VISIBLE_DEVICES
环境变量来选择要使用的 GPU。
或者,在python代码中设置:
一些stack overflow的帖子会让你尝试将这些行添加到你的代码中 ——
释放Cache
要弄清楚你的模型在 cuda 上占用了多少内存,你可以尝试:
如果你调用python的垃圾收集,并调用pytorch的清空缓存,这基本上应该让你的GPU恢复到一个干净的状态,不使用超过它需要的内存,当你开始训练下一个模型时,不必重新启动kernel。
虽然torch.cuda.empty_cache()
或gc.collect()
可以释放CUDA内存,但显然不能释放的内存返回到Python中。因此,不要把希望寄托在这些脚本上。对于JupyterLab或Colab来说,这种方式是有效的。下面是如何使用这些代码的例子:
我们马上去看一下 .detach()和.cpu():
这通常仅适用于notebooks 和 ipython
让我们看看这些之外的替代方法
使用koila python包
以下是如何使用koila的例子
如果这些都没有帮助,解决问题的唯一方法是找出:
是什么在使用内存?
与流行的看法相反,你不需要更大的GPU来训练更大的模型,你可以简单地使用梯度累积。我们稍后会讨论这个问题。我们来剖析一下错误消息,因为这通常是一个很好的提示。
这条信息告诉你总共有15.78G的GPU内存。因此,你可以检查一下的数据和模型有多大,因为你需要将它们移动到 GPU中。如果它超过了总容量,则无法在该计算机上运行,你需要将数据分块,并在 CPU 和 GPU 之间持续移动。
减小图像尺寸也有帮助,如果它说你不能使用 x MiB,因为你只有一点内存可用,找出其他进程也在使用 GPU 并释放该空间。通过运行以下命令查找 python 进程的 PID:
kill
修改模型/训练循环
现在我们进入了最后阶段。一切都试过了,没有任何东西在使用 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除以相同的数字。
累积可以通过使用回调GradientAccumulation
来完成
我们可以做的是找到一种方法来运行 32 条数据,但让它一次表现得像 64 条数据。在一般的训练循环中,在执行loss.backward()
之前,你需要将梯度归零,如果不将梯度归零,梯度将进行累积。
因此,如果你在不将梯度归零的情况下进行 2 个半批次的训练,它们的梯度回累积,最终得到以目标有效批次大小相同的梯度。在训练循环中,我们需要使用计数器根据小batchsize大小进行更新,一旦它达到预设的目标,那就是我们将梯度归零时。在那之前,他们只是通过loss.backward()
而不断积累。
英文原文:https://medium.com/@snk.nitin/how-to-solve-cuda-out-of-memory-error-850bb247cfb2