从单卡 GPU 到多卡 GPU 的 PyTorch 模型并行训练
在深度学习中，使用多个 GPU 来加速模型训练已经成为常见的做法。PyTorch 提供了方便的工具来实现单卡到多卡的转换，本文将介绍如何将一个简单的神经网络模型从单卡 GPU 变为多卡 GPU，并使用多卡并行训练来加快训练速度。
问题描述
我们将尝试解决一个简单的分类问题，使用一个包含两个全连接层的神经网络来对图像进行分类。我们的目标是将该模型从单卡 GPU 转移到多卡 GPU，并通过多卡并行训练来提高训练速度。
解决方案
单卡 GPU 实现
首先，我们定义一个简单的神经网络模型：
import torch
import torch.nn as nn
import torch.optim as optim
class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
self.fc1 = nn.Linear(784, 256)
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = torch.flatten(x, 1)
x = self.fc1(x)
x = nn.functional.relu(x)
x = self.fc2(x)
return x
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SimpleNet().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
然后，我们定义数据加载器和训练过程：
from torchvision import datasets, transforms
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('data', train=True, download=True, transform=transform),
batch_size=64, shuffle=True
)
for epoch in range(5):
for i, (images, labels) in enumerate(train_loader):
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
if i % 100 == 0:
print(f'Epoch {epoch}, Step {i}, Loss: {loss.item()}')
多卡 GPU 实现
现在，我们将模型转移到多卡 GPU 上进行训练。PyTorch 提供了 torch.nn.DataParallel 模块来实现模型的多卡并行训练：
model = nn.DataParallel(SimpleNet()).to(device)
optimizer = optim.SGD(model.parameters(), lr=0.01)
接下来，我们只需要按照之前的训练过程进行训练即可。PyTorch 会自动将数据分发到多个 GPU 上，并在反向传播时自动合并梯度。
for epoch in range(5):
for i, (images, labels) in enumerate(train_loader):
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
if i % 100 == 0:
print(f'Epoch {epoch}, Step {i}, Loss: {loss.item()}')
旅行图
journey
title PyTorch 模型从单卡 GPU 变为多卡 GPU
section 单卡 GPU 实现
code
code
section 多卡 GPU 实现
code
code
结论
通过使用 PyTorch 的 torch.nn.DataParallel 模块，我们很容易地将一个单卡 GPU 的模型转移到多卡 GPU 上进行训练。在实际应用中，可以根据实际情况调整模型和训练参数，以获得更好的训练效果和加速效果。希望本文的内容能够帮助你更好地理解如何在 PyTorch 中进行多卡 GPU 训练。
