在 PyTorch 中强制 GPU 内存限制

Question

有没有办法强制我希望可用于特定 Pytorch 实例的 GPU 内存量的最大值？例如，我的 GPU 可能有 12Gb 可用，但我想为特定进程分配最大 4Gb。

Answer 1

与会阻塞所有 CPU 内存的 tensorflow 相比，Pytorch 只使用 'it needs'。但是你可以：

减少批量大小
使用CUDA_VISIBLE_DEVICES=# of GPU（可以是倍数）来限制可以访问的GPU。

要在程序中实现此运行，请尝试：

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

Answer 2

更新（2021 年 3 月 4 日）：现在 available in the stable 1.8.0 version of PyTorch. Also, in the docs

原回答如下。

这个 feature request has been merged 进入 PyTorch master 分支。然而，未在稳定版本中引入。

Introduced as set_per_process_memory_fraction

Set memory fraction for a process. The fraction is used to limit an caching allocator to allocated memory on a CUDA device. The allowed value equals the total visible memory multiplied fraction. If trying to allocate more than the allowed value in a process, will raise an out of memory error in allocator.

您可以查看 tests 作为用法示例。

Answer 3

更新pytorch到1.8.0 （pip install --upgrade torch==1.8.0）

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

参数：

分数（浮点数） – 范围：0~1。允许的内存等于 total_memory * 分数。

设备（torch.device 或整数，可选） – 选定的设备。如果是 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)
"""

在 PyTorch 中强制 GPU 内存限制

Force GPU memory limit in PyTorch

pytorch