在PyTorch中使用GPU進行加速計算非常簡單,只需按照以下步驟操作:
首先,檢查是否有可用的GPU設備。可以通過以下代碼來獲取可用的GPU設備列表:
import torch
if torch.cuda.is_available():
device = torch.device('cuda')
print('GPU is available')
else:
device = torch.device('cpu')
print('No GPU available, using CPU instead')
接下來,將模型和數據移動到GPU設備上。可以通過以下代碼將模型移動到GPU:
model = model.to(device)
對于數據,可以通過以下代碼將Tensor移動到GPU:
data = data.to(device)
在模型和數據都移動到GPU設備后,接下來的計算都會在GPU上進行加速。在訓練時,可以通過以下代碼指定使用的設備:
model.train()
for batch in data_loader:
inputs, targets = batch[0].to(device), batch[1].to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_function(outputs, targets)
loss.backward()
optimizer.step()
通過以上步驟,就可以在PyTorch中使用GPU進行加速計算了。