运行时错误：CUDA 内存不足。如何设置 max_split_size_mb？

Question

我在 Colab Pro+ 上运行神经网络时发现这个问题（具有高 RAM 选项）。

运行时错误：CUDA 内存不足。尝试分配 8.00 GiB（GPU 0； 15.90 GiB 总容量；已分配 12.04 GiB； 2.72 GiB 免费； PyTorch 总共保留了 12.27 GiB）如果保留内存是 >> 分配的内存，请尝试设置 max_split_size_mb 以避免碎片。 请参阅内存管理和 PYTORCH_CUDA_ALLOC_CONF 的文档

我已经将批量大小减少到 2。我使用

h5py

格式上传数据。

此时，我认为我唯一可以尝试的就是设置

max_split_size_mb

。

我找不到任何关于如何实施

max_split_size_mb

的信息。我不清楚 PyTorch 文档。

Answer 1

max_split_size_mb

配置值可以设置为环境变量。

确切的语法已记录，但简而言之：

缓存分配器的行为可以通过环境变量

PYTORCH_CUDA_ALLOC_CONF

来控制。格式为

PYTORCH_CUDA_ALLOC_CONF=<option>:<value>,<option2>:<value2>

…

可用选项：

…

• max_split_size_mb

  防止分配器分割大于此大小（以 MB 为单位）的块。这有助于防止碎片，并可能允许完成一些边界工作负载而不会耗尽内存。根据分配模式，性能成本可以从“零”到“大量”不等。默认值是无限的，即所有块都可以拆分。

  memory_stats()

  和

  memory_summary()

  方法对于调整非常有用。对于由于“内存不足”而中止并显示大量不活动拆分块的工作负载，应将此选项用作最后的手段。

…

因此，您应该能够以类似于以下方式设置环境变量：

Windows：

set 'PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512'

Linux：

export 'PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512'

这取决于您使用的操作系统 - 对于您的情况，对于 Google Colab，您可能会发现 这个问题 有帮助。

Answer 2

添加到另一个答案中，要使用的大小实际上取决于错误消息中的数字，但如果您正在运行Python，

import os

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:<enter-size-here>"

在脚本开始时，我发现有时对我有用。尝试不同的尺寸。

Answer 3

尝试另一种选择：

torch.cuda.empty_cache() 

如果你想清空分配的缓存。

Answer 4

对于使用带有 GTX 1660 或其他 16XX 6GB 卡的 SD 的任何人，在使用最新的 Nvidia 驱动程序时，“实际上并不需要”此选项。当使用版本 531 时，这使得 SD 速度更快，但需要 --medvram。