PyTorch 获取 GPU 编号

PyTorch 是一款用于科学计算的开源深度学习平台。在训练神经网络时，使用 GPU 可以显著地提高运算速度。对于拥有多张 GPU 的机器来说，我们需要指定使用哪张 GPU 来进行计算。在 PyTorch 中获取 GPU 编号非常简单。

代码演示

以下是获取 GPU 编号的 PyTorch 代码：

import torch

# 判断是否可以使用 CUDA
if torch.cuda.is_available():
    # 使用 GPU 进行计算
    device = torch.device("cuda")
    print("GPU device: ", device)
else:
    # 使用 CPU 进行计算
    device = torch.device("cpu")
    print("CPU device: ", device)

上述代码会输出当前可用的设备，如果有 GPU 可用，会输出 GPU 的编号（例如 "cuda:0"），否则会输出 "cpu"。

解释说明

1. 我们首先需要导入 PyTorch 库。
2. 使用 torch.cuda.is_available() 函数判断当前是否可以使用 GPU。
3. 如果可以使用 GPU，则使用 torch.device() 函数指定使用 GPU 进行计算，并将设备名称赋值给 device 变量。
4. 如果不可以使用 GPU，则使用 CPU 进行计算，并将设备名称赋值给 device 变量。
5. 最后通过 print() 函数输出设备名称。

结论

以上就是 PyTorch 获取 GPU 编号的介绍。在实际应用中，我们需要根据具体情况选择使用哪张 GPU，以提高计算效率。