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

Question

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

原文由 vvvvv 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

您可以通过执行以下操作列出所有可用的 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>]

原文由 vvvvv 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 2

检查 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 替换为另一个数字。

原文由 Shaida Muhammad 发布，翻译遵循 CC BY-SA 4.0 许可协议