要在PyTorch中使用GPU加速計算,需要將模型和數據移動到GPU上進行處理。以下是上一節演示修改后的示例代碼,展示了如何在訓練過程中利用GPU加速計算:
import torch
import torch.nn as nn
import torch.optim as optim
import time# 檢查GPU是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 用于比較GPU對性能的提升程度,去掉注釋時使用CPU
# device = torch.device("cpu")# 定義一個簡單的神經網絡模型,并將其移動到GPU
class SimpleModel(nn.Module):def __init__(self):super(SimpleModel, self).__init__()self.fc1 = nn.Linear(60, 30)self.relu = nn.ReLU()self.fc2 = nn.Linear(30, 1)def forward(self, x):x = self.fc1(x)x = self.relu(x)x = self.fc2(x)return xmodel = SimpleModel().to(device)# 定義損失函數
criterion = nn.MSELoss()# 創建Adam優化器
optimizer = optim.Adam(model.parameters(), lr=0.01)# 準備訓練數據,并將其移動到GPU
input_data = torch.randn(100000, 60).to(device)
target_data = torch.randn(100000, 1).to(device)# 訓練模型
time_start = time.time()
for epoch in range(1000):optimizer.zero_grad()output = model(input_data)loss = criterion(output, target_data)loss.backward()optimizer.step()if epoch % 100 == 0:print(f'Epoch {epoch}, Loss: {loss.item()}')
time_stop = time.time()
print(f"time_spend = {time_stop - time_start} s")
在這個修改后的示例中,我們首先檢查GPU是否可用,并將模型和訓練數據移動到GPU設備上。通過調用.to(device)
方法,模型和數據都會被轉移到GPU上進行計算。接著,訓練過程中的計算將在GPU上加速進行,提高訓練效率。
而在我的機器上面,GPU訓練時間輸出為time_spend = 1.8002188205718994 s
,CPU訓練時間輸出為time_spend = 11.982393026351929 s