PyTorch 如何查看可用GPU？

郑玮雅 • 2024-06-05 05:29 • Python教程 • 阅读 20

简介

PyTorch 是一个流行的 Python 深度学习库，它利用图形处理单元 (GPU) 来加速训练和推理过程。在使用 PyTorch 时，了解如何查看可用的 GPU 至关重要，以便优化您的工作流程和最大限度地利用您的计算资源。本文将深入探讨如何使用 PyTorch 检查可用 GPU，并针对相关问题提供答案。

PyTorch 中查看可用 GPU 的方法

PyTorch 提供了多种方法来查看可用 GPU：

使用 torch.cuda.is_available()：

“`python
import torch

if torch.cuda.is_available():
print(“GPU 可用”)
else:
print(“GPU 不可用”)
“`

使用 torch.cuda.device_count()：

“`python
import torch

numgpus = torch.cuda.devicecount()
print(f”可用 GPU 数：{num_gpus}”)
“`

使用 torch.cuda.get_device_name()：

“`python
import torch

for i in range(torch.cuda.devicecount()):
print(f”GPU {i}：{torch.cuda.getdevice_name(i)}”)
“`

选择和使用 GPU

在查看可用 GPU 后，您需要选择要使用的 GPU。有两种方法可以执行此操作：

设置 CUDA_VISIBLE_DEVICES 环境变量：

bash export CUDA_VISIBLE_DEVICES=0

这将限制 PyTorch 仅使用指定设备索引的 GPU。

使用 torch.cuda.set_device()：

“`python
import torch

torch.cuda.set_device(0)
“`

同样，这将设置 PyTorch 使用特定设备索引的 GPU。

结论

了解如何使用 PyTorch 查看和选择可用 GPU对于优化深度学习工作流程至关重要。本文介绍了通过 torch.cuda.is_available()、torch.cuda.device_count() 和 torch.cuda.get_device_name() 检查可用 GPU 的方法。此外，我们讨论了使用 CUDA_VISIBLE_DEVICES 和 torch.cuda.set_device() 来选择和使用 GPU。

问答