在PyTorch中强制GPU内存限制

Question

有没有一种方法可以为特定Pytorch实例强制提供我想使用的GPU内存量的最大值？例如，我的GPU可能有12Gb可用，但是我想为特定进程分配最大4Gb。

Answer 1

将pytorch更新到1.8.0\n\xef\xbc\x88pip install --upgrade torch==1.8.0\xef\xbc\x89

\n

函数：torch.cuda.set_per_process_memory_fraction（分数，设备=无）

\n

参数：

\n

分数（浮点数） \xe2\x80\x93 范围：0~1。允许的内存等于总内存 * 分数。

\n

device (torch.device 或 int, 可选) \xe2\x80\x93 选择的设备。如果为 None，则使用默认 CUDA 设备。

\n

例如：

\n

import torch\ntorch.cuda.set_per_process_memory_fraction(0.5, 0)\ntorch.cuda.empty_cache()\ntotal_memory = torch.cuda.get_device_properties(0).total_memory\n# less than 0.5 will be ok:\ntmp_tensor = torch.empty(int(total_memory * 0.499), dtype=torch.int8, device=\'cuda\')\ndel tmp_tensor\ntorch.cuda.empty_cache()\n# this allocation will raise a OOM:\ntorch.empty(total_memory // 2, dtype=torch.int8, device=\'cuda\')\n\n"""\nIt raises an error as follows: \nRuntimeError: CUDA out of memory. Tried to allocate 5.59 GiB (GPU 0; 11.17 GiB total capacity; 0 bytes already allocated; 10.91 GiB free; 5.59 GiB allowed; 0 bytes reserved in total by PyTorch)\n"""\n

\n

Answer 2

更新(04-MAR-2021)：它现在在 PyTorch 的稳定 1.8.0 版本中可用。此外，在文档中

原答案如下。

此功能请求已合并到 PyTorchmaster分支。然而，没有在稳定版本中引入。

介绍为 set_per_process_memory_fraction

设置进程的内存分数。该分数用于将缓存分配器限制为在 CUDA 设备上分配的内存。允许的值等于总的可见内存乘数。如果尝试在进程中分配超过允许的值，则会在分配器中引发内存不足错误。

您可以检查测试作为使用示例。