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

爱学习的小道长2024-04-18 16:37

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