当 运行 PyTorch 模型 [具有足够的 GPU 内存] 时,为什么我的 CUDA 内存不足?

Why do I get CUDA out of memory when running PyTorch model [with enough GPU memory]?

我问这个问题是因为我在配备 8GB VRAM 的笔记本电脑上的 GTX 2070 上成功地训练了一个分段网络,并且我使用 完全相同的代码和安装的完全相同的软件库 在配备 GTX 1080TI 的台式电脑上,它仍然会出现内存不足的情况。

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

  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.

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

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

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

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

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