PyTorch检查GPU

在PyTorch中，我们可以使用以下代码片段来检查GPU是否可用以及当前使用哪个GPU：

import torch

if torch.cuda.is_available():
    device = torch.device("cuda")          # a CUDA device object
    print(f'There are {torch.cuda.device_count()} available.') #获取当前可用的GPU数量
    print(f'We will use the GPU: {torch.cuda.get_device_name(0)}') #获取当前使用的GPU名称
else:
    print("No GPU available, using CPU instead.")
    device = torch.device("cpu")

对于有多个GPU的系统，可以使用以下代码片段来检查每个GPU的使用情况：

import torch

if torch.cuda.is_available():
    device = torch.device("cuda")          # a CUDA device object
    print(f'There are {torch.cuda.device_count()} available.')
    print(f'We will use the GPU: {torch.cuda.get_device_name(0)}')
    
    for i in range(torch.cuda.device_count()):
        print(f'Device {i}: {torch.cuda.get_device_name(i)}') #获取每个GPU的名称
else:
    print("No GPU available, using CPU instead.")
    device = torch.device("cpu")

在训练神经网络时，我们还需要将张量移动到GPU上。可以使用以下代码片段将张量移动到GPU上：

import torch

if torch.cuda.is_available():
    device = torch.device("cuda")
    x = torch.randn(3, 3).to(device) #将张量移动到GPU上
    y = torch.randn(3, 3).to(device)
    z = x + y
    print(z)
else:
    print("No GPU available, using CPU instead.")

总结

PyTorch提供了简单易用的代码来检查GPU是否可用以及当前使用哪个GPU。对于多GPU系统，我们还可以检查每个GPU的使用情况。在训练神经网络时，我们还需要将张量移动到GPU上，而PyTorch也提供了方便的方法来实现此操作。