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

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

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

脚本功能概述

这段代码的主要功能是：

1. 检查PyTorch是否正确安装。
2. 检查CUDA是否可用。
3. 获取系统中可用的GPU数量，并分别输出每张GPU的详细信息。
4. 测试每张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()

代码解释

1. 检查PyTorch版本和CUDA可用性

   • torch.__version__：打印PyTorch的版本信息，确保安装正确。
   • torch.cuda.is_available()：判断CUDA是否可用。如果CUDA可用，意味着系统已正确配置了NVIDIA驱动和CUDA。

2. 获取可用GPU的数量

   • torch.cuda.device_count()：返回系统中可用的GPU数量。如果你有多张显卡，比如双卡，这里应该返回大于1的数值。

3. 输出每张GPU的名称

   • torch.cuda.get_device_name(i)：对于每张显卡，获取并输出它的名称。这样可以确认每张显卡是否被正确识别。

4. 在每张GPU上进行张量操作

   • with torch.cuda.device(i)：指定在第 i 张显卡上进行接下来的操作。
   • torch.rand(3, 3).cuda(i)：在第 i 张GPU上创建一个3x3的随机张量。这一步的成功执行证明该GPU可以正常进行计算。

5. 异常处理

   • 脚本通过 try-except 结构捕获潜在的错误。如果某张GPU无法创建张量，会抛出异常并输出错误信息，帮助用户排查问题。

使用方法

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

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

2. 运行该脚本：

   python check_gpus.py
   

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

   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')
   

   

总结

通过这个脚本，你可以轻松验证系统中的多张显卡是否被PyTorch和CUDA正确识别，并确认每张显卡是否能正常工作。该脚本在检测多GPU系统时非常实用，特别适用于需要在深度学习中使用多GPU进行并行训练的场景。

如果你在运行该脚本时遇到问题，可以检查以下几点：

1. NVIDIA驱动是否正确安装。
2. CUDA Toolkit是否与PyTorch兼容。
3. 确保你的PyTorch版本是支持CUDA的版本。

原文地址：https://blog.csdn.net/weixin_45498383/article/details/142361209

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！