如何使用 pytorch 列出所有当前可用的 GPU？

Question

我知道我可以使用 访问当前的 GPU torch.cuda.current_device()，但是如何获取所有当前可用 GPU 的列表？

Answer 1

我知道这个答案有点晚了。我认为问题的作者问的是 Pytorch 实际上不可以使用哪些设备：

• 有多少可用（可通过 获得device_count()）或
• 设备管理器句柄（可通过 获得torch.cuda.device(i)），这是其他一些答案给出的。

如果您想知道实际的 GPU 名称是什么（例如：NVIDIA 2070 GTI 等），请尝试以下操作：

import torch
for i in range(torch.cuda.device_count()):
   print(torch.cuda.get_device_properties(i).name)

注意函数的使用get_device_properties(i)。这将返回一个如下所示的对象：

_CudaDeviceProperties(name='NVIDIA GeForce RTX 2070', major=8, minor=6, total_memory=12044MB, multi_processor_count=28))

该对象包含一个名为 的属性name。您可以选择直接深入到名称属性，以获取与相关 GPU 关联的人类可读名称。

Answer 2

检查 PyTorch 有多少个 GPU 可用

import torch

num_of_gpus = torch.cuda.device_count()
print(num_of_gpus)

如果您想使用其中的第一个 GPU。

device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

如果您想使用其他 GPU，请将上述命令中的0替换为其他数字。

Answer 3

您可以通过执行以下操作列出所有可用的 GPU：

>>> import torch
>>> available_gpus = [torch.cuda.device(i) for i in range(torch.cuda.device_count())]
>>> available_gpus
[<torch.cuda.device object at 0x7f2585882b50>]