在PyTorch中使用GPU進行模型訓練可以極大地加速訓練過程。以下是一些簡單的步驟來使用GPU訓練模型:
import torch
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
model.to(device)
data.to(device)
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
# 后續的訓練步驟
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
通過這些步驟,你就可以在PyTorch中使用GPU來訓練模型了。使用GPU訓練模型可以顯著提高訓練速度和效率。