如何检查PyTorch是否使用GPU?nvidia-smi命令可以检测GPU活动,但我想直接从Python脚本中检查它。


当前回答

查询是否有可用的GPU。

torch.cuda.is_available()

如果上面的函数返回False,

你要么没有GPU, 或者没有安装Nvidia驱动程序,所以OS看不到GPU, 或者GPU被环境变量CUDA_VISIBLE_DEVICES隐藏。当CUDA_VISIBLE_DEVICES的值为-1时,将隐藏所有设备。你可以用下面这行代码检查这个值:os.environ['CUDA_VISIBLE_DEVICES']

如果上面的函数返回True,这并不一定意味着你正在使用GPU。在Pytorch中,您可以在创建设备时将张量分配给它们。默认情况下,张量被分配给cpu。要检查张量的分配位置,请执行以下操作:

# assuming that 'a' is a tensor created somewhere else
a.device  # returns the device where the tensor is allocated

注意,您不能操作在不同设备中分配的张量。要了解如何将张量分配给GPU,请参见这里:https://pytorch.org/docs/stable/notes/cuda.html

其他回答

Query Command
Does PyTorch see any GPUs? torch.cuda.is_available()
Are tensors stored on GPU by default? torch.rand(10).device
Set default tensor type to CUDA: torch.set_default_tensor_type(torch.cuda.FloatTensor)
Is this tensor a GPU tensor? my_tensor.is_cuda
Is this model stored on the GPU? all(p.is_cuda for p in my_model.parameters())

从实际的角度来看,有一个小题外话:

import torch
dev = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

这个开发人员现在知道是cuda还是cpu。

在使用cuda时,处理模型和张量的方式是不同的。一开始有点奇怪。

import torch
import torch.nn as nn
dev = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
t1 = torch.randn(1,2)
t2 = torch.randn(1,2).to(dev)
print(t1)  # tensor([[-0.2678,  1.9252]])
print(t2)  # tensor([[ 0.5117, -3.6247]], device='cuda:0')
t1.to(dev)
print(t1)  # tensor([[-0.2678,  1.9252]])
print(t1.is_cuda) # False
t1 = t1.to(dev)
print(t1)  # tensor([[-0.2678,  1.9252]], device='cuda:0')
print(t1.is_cuda) # True

class M(nn.Module):
    def __init__(self):        
        super().__init__()        
        self.l1 = nn.Linear(1,2)

    def forward(self, x):                      
        x = self.l1(x)
        return x
model = M()   # not on cuda
model.to(dev) # is on cuda (all parameters)
print(next(model.parameters()).is_cuda) # True

这一切都很棘手,一旦理解它,就可以帮助您快速处理较少的调试。

简单地从命令提示符或Linux环境中运行以下命令。

python -c 'import torch; print(torch.cuda.is_available())'

上面应该打印为True

python -c 'import torch; print(torch.rand(2,3).cuda())'

它应该打印以下内容:

tensor([[0.7997, 0.6170, 0.7042], [0.4174, 0.1494, 0.0516]], device='cuda:0')

查询是否有可用的GPU。

torch.cuda.is_available()

如果上面的函数返回False,

你要么没有GPU, 或者没有安装Nvidia驱动程序,所以OS看不到GPU, 或者GPU被环境变量CUDA_VISIBLE_DEVICES隐藏。当CUDA_VISIBLE_DEVICES的值为-1时,将隐藏所有设备。你可以用下面这行代码检查这个值:os.environ['CUDA_VISIBLE_DEVICES']

如果上面的函数返回True,这并不一定意味着你正在使用GPU。在Pytorch中,您可以在创建设备时将张量分配给它们。默认情况下,张量被分配给cpu。要检查张量的分配位置,请执行以下操作:

# assuming that 'a' is a tensor created somewhere else
a.device  # returns the device where the tensor is allocated

注意,您不能操作在不同设备中分配的张量。要了解如何将张量分配给GPU,请参见这里:https://pytorch.org/docs/stable/notes/cuda.html

这里几乎所有的答案都引用torch.cuda.is_available()。然而,这只是硬币的一部分。它告诉你GPU(实际上是CUDA)是否可用,而不是它是否实际被使用。在一个典型的设置中,你会像这样设置你的设备:

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

但在更大的环境中(例如研究),通常也会给用户更多的选择,因此根据输入,他们可以禁用CUDA,指定CUDA id,等等。在这种情况下,是否使用GPU不仅仅取决于GPU是否可用。在设备被设置为torch设备之后,您可以获取它的type属性来验证它是否是CUDA。

if device.type == 'cuda':
    # do something