学堂 学堂 学堂公众号手机端

PyTorch中怎么使用GPU进行计算

lewis 3年前 (2022-01-21) 阅读数 5 #技术

在PyTorch中使用GPU进行计算非常简单,只需将模型和数据加载到GPU上即可。

  1. 检查GPU是否可用:
import torch
if torch.cuda.is_available():
    device = torch.device("cuda")
    print("GPU is available")
else:
    device = torch.device("cpu")
    print("GPU is not available, using CPU")
  1. 将模型加载到GPU上:
model = Model()
model.to(device)
  1. 将数据加载到GPU上:
data = data.to(device)
  1. 在训练循环中,确保每个batch的数据都加载到GPU上:
for batch in dataloader:
    inputs, labels = batch
    inputs, labels = inputs.to(device), labels.to(device)
    # 在GPU上进行计算

通过以上步骤,就可以在PyTorch中轻松地使用GPU进行计算。


版权声明

本文仅代表作者观点,不代表博信信息网立场。

热门