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

功能：

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.org/docs/1.8.0/cuda.html

示例：

import