【技巧】PyTorch限制GPU显存的可使用上限

转载请注明出处：小锋学长生活大爆炸[xfxuezhang.cn]

        从 PyTorch 1.4 版本开始，引入了一个新的功能 torch.cuda.set_per_process_memory_fraction(fraction, device)，这个功能允许用户为特定的 GPU 设备设置进程可使用的显存上限比例。

        测试代码：

torch.cuda.empty_cache()
 
# 设置进程可使用的GPU显存最大比例为50%
torch.cuda.set_per_process_memory_fraction(0.5, device=0)
 
# 计算总内存
total_memory = torch.cuda.get_device_properties(0).total_memory
print("实际总内存:", round(total_memory / (1024 * 1024), 1), "MB")
 
# 尝试分配大量显存的操作
try:
    # 使用10%的显存:
    tmp_tensor = torch.empty(int(total_memory * 0.1), dtype=torch.int8, device='cuda:0')
    print("分配的内存:", round(torch.cuda.memory_allocated(0) / (1024 * 1024), 1), "MB")
    print("保留的内存:", round(torch.cuda.memory_reserved(0) / (1024 * 1024), 1), "MB")
    # 清空显存
    del tmp_tensor
    torch.cuda.empty_cache()
    # 使用50%的显存:
    torch.empty(int(total_memory * 0.5), dtype=torch.int8, device='cuda:0')
except RuntimeError as e:
    print("Error allocating tensor:", e)
 
# 打印当前GPU的显存使用情况
print("分配的内存:", torch.cuda.memory_allocated(0) / (1024 * 1024), "MB")
print("保留的内存:", torch.cuda.memory_reserved(0) / (1024 * 1024), "MB")

• 已分配显存：通过torch.cuda.memory_allocated(device)查询，它返回已经直接分配给张量的显存总量。这部分显存是当前正在被Tensor对象使用的。

• 保留（预留）显存：通过torch.cuda.memory_reserved(device)查询，它包括了已分配显存以及一部分由PyTorch的CUDA内存分配器为了提高分配效率和减少CUDA操作所需时间而预留的显存。这部分预留的显存不直接用于存储Tensor对象的数据，但可以被视为快速响应未来显存分配请求的“缓冲区”。