pytorch查看cuda显存

本文介绍如何使用pytorch来列出当前所有可用的GPU以及cuda相关信息；

检查CUDA和PyTorch的安装

需要确保我们已经正确地安装了cuda和pytorch，cuda是NVIDIA提供的用于GPU加速计算的平台，pytorch是基于CUDA开发的。

使用以下代码片段来检查cuda的安装情况：

import torch
print(torch.cuda.is_available())

输出结果为True，则表示CUDA已经安装成功；

import torch
print(torch.__version__)

输出结果为PyTorch的版本号，说明PyTorch已经正确地安装；

示例：

列出当前可用的GPU

pytorch提供了一个torch.cuda.device_count()方法来获取当前可用的GPU数量；

import torch
num_gpus = torch.cuda.device_count()
print("当前可用的GPU数量：", num_gpus)

示例：

获取GPU的详细信息

pytorch可以获取更详细的GPU信息，例如GPU的名称、显存大小、已使用的gpu显存、剩余显存等；

import torch
num_gpus = torch.cuda.device_count()
for i in range(num_gpus):
    device = torch.device(f"cuda:{i}")
    properties = torch.cuda.get_device_properties(device)
    total_memory = properties.total_memory / (1024 ** 3)  # 显存总量(GB)
    used_memory = torch.cuda.memory_allocated(torch.cuda.current_device()) / (1024 ** 3)  # 已使用显存(GB)
    free_memory = total_memory - used_memory  # 剩余显存(GB)
    print(f"GPU {i} 的详细信息：")
    print("名称：", properties.name)
    print(f"显存大小：{total_memory:.2f} GB”)
    print(f"已使用的GPU显存：{used_memory:.2f} GB")
    print(f"剩余GPU显存：{free_memory:.2f} GB")

示例：

设置使用的GPU

在有多个GPU的情况下，我们可以通过设置torch.cuda.set_device()方法来选择使用哪个GPU；

import torch
torch.cuda.set_device(0)