[Python] pytorch, whether CUDA is available, check the remaining capacity of the graphics card memory

CUDA is available, and a total of 1 GPU device is available.
Currently used GPU device index: 0
Currently used GPU device name: NVIDIA T1000
Total GPU memory: 4.00 GB
Used GPU memory: 0.00 GB
Remaining GPU memory: 4.00 GB
PyTorch version: 1.10.1+cu102

import torch

# 检查CUDA是否可用
cuda_available = torch.cuda.is_available()

if cuda_available:
    # 获取GPU设备数量
    num_gpu = torch.cuda.device_count()

    # 获取当前使用的GPU索引
    current_gpu_index = torch.cuda.current_device()

    # 获取当前GPU的名称
    current_gpu_name = torch.cuda.get_device_name(current_gpu_index)

    # 获取GPU显存的总量和已使用量
    total_memory = torch.cuda.get_device_properties(current_gpu_index).total_memory / (1024 ** 3)  # 显存总量(GB)
    used_memory = torch.cuda.memory_allocated(current_gpu_index) / (1024 ** 3)  # 已使用显存(GB)
    free_memory = total_memory - used_memory  # 剩余显存(GB)

    print(f"CUDA可用，共有 {
      
      num_gpu} 个GPU设备可用。")
    print(f"当前使用的GPU设备索引：{
      
      current_gpu_index}")
    print(f"当前使用的GPU设备名称：{
      
      current_gpu_name}")
    print(f"GPU显存总量：{
      
      total_memory:.2f} GB")
    print(f"已使用的GPU显存：{
      
      used_memory:.2f} GB")
    print(f"剩余GPU显存：{
      
      free_memory:.2f} GB")
else:
    print("CUDA不可用。")

# 检查PyTorch版本
print(f"PyTorch版本：{
      
      torch.__version__}")

Windows installs the graphics card driver first, then CUDA10.2, and finally pytorch.

pip install torch1.10.1+cu102 torchvision0.13.1+cu102 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu102