【技巧】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对象的数据，但可以被视为快速响应未来显存分配请求的“缓冲区”。
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。
评论
登录后参与评论
推荐阅读