常用代码——Pytorch进行显存预先占用_如何预先占用测试的显存

import torch
# 限制0号设备的显存的使用量为0.5，就是半张卡那么多，比如12G卡，设置0.5就是6G。
torch.cuda.set_per_process_memory_fraction(0.5, 0)
torch.cuda.empty_cache()
# 计算一下总内存有多少。
total_memory = torch.cuda.get_device_properties(0).total_memory
# 使用0.499的显存:
tmp_tensor = torch.empty(int(total_memory * 0.499), dtype=torch.int8, device='cuda')

# 清空该显存：
del tmp_tensor
torch.cuda.empty_cache()

# 下面这句话会触发显存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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

import torch
# 限制0号设备的显存的使用量为1.0，就是半张卡那么多，比如12G卡，设置1.0就是12G。
torch.cuda.set_per_process_memory_fraction(1.0, 0)
torch.cuda.empty_cache()
# 计算一下总内存有多少。
total_memory = torch.cuda.get_device_properties(0).total_memory
# 使用0.9的显存:
tmp_tensor = torch.empty(int(total_memory * 0.9), dtype=torch.int8, device='cuda
1
2
3
4
5
6
7
8