强制限制PyTorch中的GPU内存使用

Question

有没有一种方式可以强制限制要分配给特定Pytorch实例的GPU内存的最大值？例如，我的GPU可能有12GB可用，但我希望将最大值分配为4GB，以供特定进程使用。

- Giorgos Sfikas

---

我认为PyTorch会根据需要使用内存，可能包括模型和已加载的图像。我不认为它会贪婪地占用所有可用内存。 - Manuel Lagunas

@ManuelLagunas 这是这样吗？即使如此，我还想知道是否有一种方法可以设置 pytorch 实例可见的内存的最大限制。 - Giorgos Sfikas

我其实不知道具体的方法，但我认为你可以通过批量大小来控制。模型大小始终保持不变，如果它不能适应GPU内存，执行就会崩溃。 - Manuel Lagunas

Answer 1

更新（2021年3月4日）：现在PyTorch稳定版1.8.0版本中可用。此外，在文档中也有提到。

原始答案如下。

---

这个功能请求已经合并到了PyTorch的master分支中，但还没有在稳定版本中引入。 作为set_per_process_memory_fraction引入

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

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

- ndrwnaguib

Answer 2

更新 PyTorch 到 1.8.0 版本（使用命令：pip install --upgrade torch==1.8.0）

函数：torch.cuda.set_per_process_memory_fraction(fraction, device=None)

参数：

fraction (float) – 取值范围为 0~1，表示允许使用的显存大小为总显存大小乘以 fraction。

device (torch.device 或 int, 可选) – 指定设备。如果为 None，则使用默认的 CUDA 设备。

示例：

import torch
torch.cuda.set_per_process_memory_fraction(0.5, 0)
torch.cuda.empty_cache()
total_memory = torch.cuda.get_device_properties(0).total_memory
# less than 0.5 will be ok:
tmp_tensor = torch.empty(int(total_memory * 0.499), dtype=torch.int8, device='cuda')
del tmp_tensor
torch.cuda.empty_cache()
# this allocation will raise a OOM:
torch.empty(total_memory // 2, dtype=torch.int8, device='cuda')

"""
It raises an error as follows: 
RuntimeError: 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)
"""

- kaiyuanxie

Answer 3

与TensorFlow不同，PyTorch仅使用其所需的内存，而不会阻塞所有CPU内存。但是，您可以：

• 减少批量大小
• 使用CUDA_VISIBLE_DEVICES = ＃个GPU（可以是多个）来限制可以访问的GPU。

为了使其在程序中运行，请尝试：

import os
os.environ["CUDA_VISIBLE_DEVICES"]="0"

- Benedict K.

---

这不是一个真正的答案。以上两种方法都会影响内存使用量，但都不会限制 Pytorch。我发现这是 Pytorch 的一个相当恼人的特性，因为内存分配允许更好地共享资源（是的，有时以减慢运行时间为代价）。 此外，还有一个官方的功能请求： https://github.com/pytorch/pytorch/issues/18626 - Ido_f

实际上，这个答案并没有解决如何强制限制内存使用的问题。此外，pytorch并不只是保留所需GPU内存。Pytorch将不再使用的GPU内存（例如一个张量变量超出范围）保留给未来的分配，而不是释放给操作系统。这意味着，即使两个进程在任何特定时间使用的GPU内存总和仍低于容量，它们仍然会遇到内存不足错误。 - maarten