在Torch中實現多GPU訓練可以通過使用nn.DataParallel
模塊來實現。nn.DataParallel
模塊可以將模型復制到多個GPU上,并自動將數據分發到每個GPU上進行計算,并最終將結果進行合并。以下是一個簡單示例:
import torch
import torch.nn as nn
# 定義模型
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 1)
def forward(self, x):
x = self.fc1(x)
x = self.fc2(x)
return x
model = Model()
# 檢查是否有多個GPU可用
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
# 將模型加載到GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# 定義優化器和損失函數
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
criterion = nn.MSELoss()
# 訓練模型
for epoch in range(num_epochs):
for inputs, targets in dataloader:
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
在上面的示例中,首先檢查是否有多個GPU可用,如果有,則將模型放入nn.DataParallel
中。然后將模型加載到GPU上進行訓練。在訓練過程中,數據會自動分發到每個GPU上,并在多個GPU上并行計算,從而加快訓練速度。