首页 > 解决方案 > 为什么在运行 PyTorch 模型 [具有足够 GPU 内存] 时 CUDA 内存不足?

问题描述

我问这个问题是因为我成功地在具有 8GB VRAM 的笔记本电脑上的 GTX 2070 上训练了一个分段网络,并且我使用完全相同的代码和完全相同的软件库安装在带有 GTX 1080TI 的台式 PC 上,它仍然抛出记忆。

为什么会发生这种情况,考虑到:

  1. 相同的 Windows 10 + CUDA 10.1 + CUDNN 7.6.5.32 + Nvidia 驱动程序 418.96(与 CUDA 10.1 一起提供)在笔记本电脑和 PC 上都有。

  2. 事实上,使用 TensorFlow 2.3 进行的训练在我 PC 上的 GPU 上运行流畅,但它无法为仅使用 PyTorch 的训练分配内存。

  3. PyTorch 通过以下命令识别 GPU(打印 GTX 1080 TI):print(torch.cuda.get_device_name(0))

  4. PyTorch 在运行此命令时分配内存:torch.rand(20000, 20000).cuda() #allocated 1.5GB of VRAM.

解决方案是什么?

标签: pythontensorflowpytorchout-of-memory

解决方案


大多数人(甚至在下面的线程中)都跳到建议减少 batch_size 将解决这个问题。事实上,在这种情况下并非如此。例如,考虑到在具有 11GB VRAM 的系统上没有其他应用程序消耗视频内存并且安装了完全相同的配置和用过的。

在我的情况下发生这种情况的原因是,在使用DataLoader对象时,我为参数设置了一个非常高的 (12) 值workers。在我的情况下,将此值减小到 4 解决了这个问题。

事实上,尽管在线程的底部,Yurasyk 在https://github.com/pytorch/pytorch/issues/16417#issuecomment-599137646提供的答案为我指明了正确的方向。

解决方案:减少workersPyTorch 中的数量DataLoader。虽然我不完全理解为什么这个解决方案有效,但我认为它与在幕后产生的用于数据获取的线程有关;在某些处理器上,可能会出现这样的错误。


推荐阅读