在PyTorch中使用GPU進行訓練非常簡單,只需按照以下步驟操作:
import torch
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
model = YourModel().to(device)
data = YourDataLoader().to(device)
for inputs, labels in data:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
# 后續訓練步驟
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
optimizer.zero_grad()
loss.backward()
optimizer.step()
通過上述步驟,就可以很方便地在PyTorch中使用GPU進行模型訓練。