功能：

pytorch（版本1.8.0以后）中设置显存使用比例，即能够完成显存的使用上限设置，用户可以根据需要限制GPU的显存消耗上限。

torch.cuda.set_per_process_memory_fraction(0.5, 0)

     参数1：fraction 限制的上限比例，如0.5 就是总GPU显存的一半，可以是0~1的任意float大小；
     参数2：device 设备号； 如0 表示GPU卡 0号；

功能解释如下：
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.
参看https://pytorch/docs/1.8.0/cuda.html

示例：

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

多进程的限制：

该函数只能限制进程的显存量，如果需要限制子进程，需要进入到子进程中限制。如下所示：

import torch
from torch.multiprocessing import Process

def sub_func(fraction, device):
    torch.cuda.set_per_process_memory_fraction(fraction, device)
    pass #do something

if __name__ == "__main__":
    for _ in range(2):
        # 每个子进程在卡0上占用0.3的显存。
        p = Process(target=main, args=(0.3, 0))
        p.start()
        ps.append(p)
    for p in ps:
        p.join()

更多推荐

pytorch的torch.cuda.set_per_process_memory_fraction()函数介绍