如何检查 PyTorch 是否正在使用 GPU？

Question

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

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

Answer 1

这些功能应该有助于：

 >>> import torch

>>> torch.cuda.is_available()
True

>>> torch.cuda.device_count()
1

>>> torch.cuda.current_device()
0

>>> torch.cuda.device(0)
<torch.cuda.device at 0x7efce0b03be0>

>>> torch.cuda.get_device_name(0)
'GeForce GTX 950M'

这告诉我们：

• CUDA 可用并且可由一台设备使用。
• Device 0 指的是GPU GeForce GTX 950M 目前被PyTorch选中。

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

Answer 2

由于这里没有提出，我添加了一个使用 torch.device 的方法，因为这非常方便，在正确的 device 上初始化张量时也是如此。

 # setting device on GPU if available, else CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
print()

#Additional Info when using cuda
if device.type == 'cuda':
    print(torch.cuda.get_device_name(0))
    print('Memory Usage:')
    print('Allocated:', round(torch.cuda.memory_allocated(0)/1024**3,1), 'GB')
    print('Cached:   ', round(torch.cuda.memory_reserved(0)/1024**3,1), 'GB')

编辑： torch.cuda.memory_cached 已重命名为 torch.cuda.memory_reserved 。因此，对于旧版本，请使用 memory_cached 。

输出：

 Using device: cuda

Tesla K80
Memory Usage:
Allocated: 0.3 GB
Cached:    0.6 GB

如上所述，使用 device _可以_：

• 将 张量移动到各自的 device ：

   torch.rand(10).to(device)

• 直接在 device 上 创建 张量：

   torch.rand(10, device=device)

这使得在 CPU 和 GPU 之间切换很舒服，而无需更改实际代码。

---

编辑：

由于对 缓存 和 分配 的内存存在一些疑问和困惑，我添加了一些关于它的附加信息：

• torch.cuda.max_memory_cached(device=None)

返回给定设备的缓存分配器管理的最大 GPU 内存（以字节为单位）。

• torch.cuda.memory_allocated(device=None)

返回给定设备的张量当前 GPU 内存使用情况（以字节为单位）。

您可以直接交出 device ，如上文所述，也可以保留 None ，它将使用 current_device() 。

---

附加说明：具有 Cuda 计算能力 3.0 或更低版本的旧图形卡可能是可见的，但不能被 Pytorch 使用！

感谢 hekimgil 指出这一点！ - “发现 GPU0 GeForce GT 750M 具有 cuda 功能 3.0。PyTorch 不再支持此 GPU，因为它太旧了。我们支持的最低 cuda 功能是 3.5。”

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