Pytorch 获取当前模型占用的 GPU显存的大小

本文主要是介绍Pytorch 获取当前模型占用的 GPU显存的大小，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

1. 背景

AI 训练模型时需要知道模型占用显存的大小，以及剩余显存大小

2. 代码

import torch# 检查GPU是否可用
if torch.cuda.is_available():# 获取当前设备device = torch.cuda.current_device()# 创建一个虚拟modelmodel = torch.nn.Linear(25600, 20000, 20000).cuda()# 获取当前模型占用的显存大小model_memory = torch.cuda.memory_allocated()# 获取PyTorch预留的显存大小pytorch_memory = torch.cuda.max_memory_allocated() - model_memory# 打印模型占用的显存大小和PyTorch预留的显存大小print(f"Model memory usage: {model_memory / (1024 ** 3):.2f} GB")print(f"PyTorch memory usage: {pytorch_memory / (1024 ** 3):.2f} GB")# 获取已使用的显存大小total_memory = torch.cuda.get_device_properties(device=device).total_memory# 获取剩余的显存大小free_memory = total_memory - model_memoryprint(f"剩余的显存大小:{free_memory / (1024 ** 3):.2f} GB")# 释放虚拟张量占用的显存del modeltorch.cuda.empty_cache()
else:print("GPU is not available.")

运行结果：

Model memory usage: 1.91 GB
PyTorch memory usage: 0.00 GB
剩余的显存大小:3.89 GB

这篇关于Pytorch 获取当前模型占用的 GPU显存的大小的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！

---