memory - RuntimeError : CUDA out of memory. 如何设置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 文档 ( https://pytorch.org/docs/stable/notes/cuda.html ) 对我来说不是很清楚。

有人可以支持我吗？谢谢。

Answer 1

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

确切的语法记录在 https://pytorch.org/docs/stable/notes/cuda.html#memory-management ，但简而言之:

The behavior of caching allocator can be controlled via environment variable PYTORCH_CUDA_ALLOC_CONF. The format is PYTORCH_CUDA_ALLOC_CONF=<option>:<value>,<option2>:<value2>...

Available options:

• max_split_size_mb prevents the allocator from splitting blocks larger than this size (in MB). This can help prevent fragmentation and may allow some borderline workloads to complete without running out of memory. Performance cost can range from ‘zero’ to ‘substatial’ depending on allocation patterns. Default value is unlimited, i.e. all blocks can be split. The memory_stats() and memory_summary() methods are useful for tuning. This option should be used as a last resort for a workload that is aborting due to ‘out of memory’ and showing a large amount of inactive split blocks.

...

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

window :set 'PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512'

Linux:export 'PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512'

这将取决于您使用的操作系统 - 在您的情况下，对于 Google Colab，您可能会找到 Setting environment variables in Google Colab有帮助。