使用PyTorch检测和验证多GPU环境的Python脚本

news2024/9/23 18:34:49

使用PyTorch检测和验证多GPU环境的Python脚本

在深度学习和机器学习中，GPU的计算能力对模型训练和推理的速度有着极大的影响。随着多GPU系统的普及，如何确保多GPU能被正确识别并使用，是一个非常关键的问题。本文将为大家介绍一段简洁的Python脚本，它可以帮助你验证PyTorch安装是否正常，并确保系统中的多张GPU都能被正确识别和使用。

脚本功能概述

这段代码的主要功能是：

检查PyTorch是否正确安装。
检查CUDA是否可用。
获取系统中可用的GPU数量，并分别输出每张GPU的详细信息。
测试每张GPU是否能正常工作，在GPU上创建张量。

脚本代码

import torch

def check_multiple_gpus():
    # 检查PyTorch是否安装以及版本信息
    print(f"PyTorch Version: {torch.__version__}")

    # 检查是否安装了CUDA
    cuda_available = torch.cuda.is_available()
    print(f"CUDA Available: {cuda_available}")

    if cuda_available:
        # 获取可用GPU的数量
        gpu_count = torch.cuda.device_count()
        print(f"Number of GPUs available: {gpu_count}")

        # 循环遍历每一张GPU
        for i in range(gpu_count):
            print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
            
            # 测试在该GPU上创建张量
            try:
                with torch.cuda.device(i):
                    x = torch.rand(3, 3).cuda(i)
                    print(f"Tensor successfully created on GPU {i}: \n{x}")
            except Exception as e:
                print(f"Error when creating tensor on GPU {i}: {e}")
    else:
        print("CUDA is not available, make sure it is installed correctly.")

if __name__ == "__main__":
    check_multiple_gpus()

代码解释

检查PyTorch版本和CUDA可用性
- torch.__version__：打印PyTorch的版本信息，确保安装正确。
- torch.cuda.is_available()：判断CUDA是否可用。如果CUDA可用，意味着系统已正确配置了NVIDIA驱动和CUDA。
获取可用GPU的数量
- torch.cuda.device_count()：返回系统中可用的GPU数量。如果你有多张显卡，比如双卡，这里应该返回大于1的数值。
输出每张GPU的名称
- torch.cuda.get_device_name(i)：对于每张显卡，获取并输出它的名称。这样可以确认每张显卡是否被正确识别。
在每张GPU上进行张量操作
- with torch.cuda.device(i)：指定在第 i 张显卡上进行接下来的操作。
- torch.rand(3, 3).cuda(i)：在第 i 张GPU上创建一个3x3的随机张量。这一步的成功执行证明该GPU可以正常进行计算。
异常处理
- 脚本通过 try-except 结构捕获潜在的错误。如果某张GPU无法创建张量，会抛出异常并输出错误信息，帮助用户排查问题。

使用方法

确保你已经正确安装了PyTorch和CUDA驱动。可以通过以下命令安装PyTorch：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

运行该脚本：
```
python check_gpus.py
```

脚本将输出类似如下的信息：

PyTorch Version: 2.0.1
CUDA Available: True
Number of GPUs available: 2
GPU 0: NVIDIA GeForce RTX 4090 D
Tensor successfully created on GPU 0: 
tensor([[0.1409, 0.2232, 0.2413],
        [0.5565, 0.3695, 0.5665],
        [0.3443, 0.6117, 0.1937]], device='cuda:0')
GPU 1: NVIDIA GeForce RTX 4090 D
Tensor successfully created on GPU 1: 
tensor([[0.5749, 0.4472, 0.0666],
        [0.7913, 0.7882, 0.8210],
        [0.9765, 0.4896, 0.9297]], device='cuda:1')