【Pytorch】GPU内存爆炸的原因

最新推荐文章于 2023-12-13 10:28:59 发布

原创最新推荐文章于 2023-12-13 10:28:59 发布 · 2.3k 阅读

12 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #人工智能 #机器学习 #python

Pytorch 专栏收录该内容

4 篇文章

订阅专栏

本文记录了在运行深度学习模型时遇到的CUDA out of memory问题，分析了可能的原因，包括测试阶段的梯度计算、其他进程占用GPU资源等，并提供了相应的解决策略，如使用`torch.no_grad()`上下文管理器、杀掉占用进程以及在训练和测试后释放GPU内存的代码示例。同时，文章鼓励读者针对类似问题进行深入研究和讨论。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

【1】出现的问题：

RuntimeError: CUDA out of memory.

也就是常说的显存爆炸。

查了一下大致可能出现这个问题的原因：

（1）在测试阶段使用了梯度计算，方法就是测试阶段添加以下代码

with torch.no_grad():#测试阶段

（2）有其它进程占用，方法就是杀进程

taskkill /PID 进程号 -t -f

（3）就是需要释放GPU的内存，在训练和测试阶段均添加下面的代码

for batch, (X, y) in enumerate(dataloader):
    X = X.to(device)
    y = y.to(device)
    pred = model(X)
    loss = loss_fn(pred, y)

    ······（省略了一些代码）

    # 添加如下代码
    del X, y, pred, loss
    torch.cuda.empty_cache()

以上大概就是常见的三种错误，当然还有其它的错误，这里就不在列举，可以去其它地方找找看。

以上就是今天的记录，欢迎大家提出问题！