在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進行計算。