在PyTorch中使用GPU进行计算非常简单,只需将模型和数据加载到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")
model = Model()
model.to(device)
data = data.to(device)
for batch in dataloader:
inputs, labels = batch
inputs, labels = inputs.to(device), labels.to(device)
# 在GPU上进行计算
通过以上步骤,就可以在PyTorch中轻松地使用GPU进行计算。