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

1. 背景

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

2. 代码

bash 复制代码
import torch


# 检查GPU是否可用
if torch.cuda.is_available():
    # 获取当前设备
    device = torch.cuda.current_device()
    # 创建一个虚拟model
    model = 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_memory
    print(f"剩余的显存大小:{free_memory / (1024 ** 3):.2f} GB")

    # 释放虚拟张量占用的显存
    del model
    torch.cuda.empty_cache()
else:
    print("GPU is not available.")

运行结果:

bash 复制代码
Model memory usage: 1.91 GB
PyTorch memory usage: 0.00 GB
剩余的显存大小:3.89 GB
相关推荐
youcans_42 分钟前
2025年数学建模美赛 A题分析(4)楼梯使用人数模型
python·数学建模
PPT百科44 分钟前
创建实用PPT演讲者备注的有效方法
人工智能·经验分享·pdf·powerpoint·ppt
查理零世3 小时前
【算法】数论基础——约数个数定理、约数和定理 python
python·算法·数论
lilu88888884 小时前
AI代码生成器赋能房地产:ScriptEcho如何革新VR/AR房产浏览体验
前端·人工智能·ar·vr
Eiceblue4 小时前
Python 合并 Excel 单元格
开发语言·vscode·python·pycharm·excel
梦云澜4 小时前
论文阅读(五):乳腺癌中的高斯图模型和扩展网络推理
论文阅读·人工智能·深度学习·学习
危险、5 小时前
Spring Boot 无缝集成SpringAI的函数调用模块
人工智能·spring boot·函数调用·springai
深度学习实战训练营5 小时前
基于迁移学习的ResNet50模型实现石榴病害数据集多分类图片预测
人工智能·分类·迁移学习
XianxinMao6 小时前
开源AI模型发布策略:平衡开放与质量的艺术
人工智能
Fxrain6 小时前
[Computer Vision]实验二:图像特征点提取
人工智能·计算机视觉