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

Posted 2023-03-12

【中文标题】如何使用 pytorch 列出所有当前可用的 GPU？

【英文标题】：How do I list all currently available GPUs with pytorch?

【发布时间】：2021-02-22 20:36:40

【问题描述】：

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

【参考方案1】：

检查 PyTorch 有多少 GPU 可用

import torch

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

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

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

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

【参考方案2】：

您可以通过以下方式列出所有可用的 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>]

【讨论】：

这不是一个正确的答案。 torch.cuda.device(i) 返回一个上下文管理器，它会导致未来的命令使用该设备。将它们全部放在这样的列表中是没有意义的。你真正需要的是torch.cuda.device_count()，你的cuda设备是cuda:0、cuda:1等，直到device_count() - 1。