PyTorch 检查是否使用 GPU - Python

当使用 PyTorch 训练深度学习模型时，通常会使用 GPU，因为与 CPU 相比，GPU 可以更快地执行矩阵计算和深度学习算法。

在本文中，我们将介绍如何在 PyTorch 中检查是否使用了 GPU，以确保我们的深度学习模型正在使用 GPU 进行训练。

检查是否使用 GPU

要检查 PyTorch 是否在 GPU 上运行，请使用以下代码：

import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)

如果 PyTorch 正在使用 GPU 运行，则上述代码将输出 Using device: cuda。如果 PyTorch 正在使用 CPU 运行，则输出 Using device: cpu。

将模型移动到 GPU 上

如果我们发现 PyTorch 在 CPU 上运行，我们可以将模型移动到 GPU 上，以加速深度学习模型的训练。我们可以使用以下代码将模型移动到 GPU 上：

model.to(device)

使用 model.to(device) 将所有模型参数移动到 GPU 上。

结论

在本文中，我们介绍了如何在 PyTorch 中检查是否使用 GPU，并将深度学习模型移动到 GPU 上进行训练。通过在 GPU 上进行训练，可以大大加速深度学习模型的训练过程，同时提高深度学习模型的预测性能。