使用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 =</