为什么在运行 PyTorch 模型时 [使用足够的 GPU 内存] 会导致 CUDA 内存不足?

Tim*_*lin 6 python out-of-memory tensorflow pytorch

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

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

  1. 相同的 Windows 10 + CUDA 10.1 + CUDNN 7.6.5.32 + Nvidia Driver 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.

解决这个问题的方法是什么?

Tim*_*lin 9

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

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

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

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