功能:
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