运行时错误：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值可以设置为环境变量。

\n

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

\n

\n

缓存分配器的行为可以通过环境变量控制PYTORCH_CUDA_ALLOC_CONF。格式为PYTORCH_CUDA_ALLOC_CONF=<option>:<value>,<option2>:<value2>\xe2\x80\xa6

\n

可用选项：

\n

\xe2\x80\xa6

\n

\n
• max_split_size_mb防止分配器分割大于此大小（以 MB 为单位）的块。这有助于防止碎片，并可能允许完成一些边界工作负载而不会耗尽内存。性能成本的范围可以从 \xe2\x80\x98zero\xe2\x80\x99 到 \xe2\x80\x98substantial\xe2\x80\x99，具体取决于分配模式。默认值是无限制的，即所有块都可以拆分。和memory_stats()方法memory_summary()对于调整很有用。对于由于 \xe2\x80\x98 内存不足\xe2\x80\x99 并显示大量非活动分割块而中止的工作负载，应将此选项用作最后的手段。
\n

\n

\xe2\x80\xa6

\n

\n

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

\n

视窗：set \'PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512\'

\n

Linux：export \'PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512\'

\n

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

\n

Answer 2

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

import os

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

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