PyTorch 多 GPU 训练教程

PyTorch 是一个流行的深度学习框架，支持多 GPU 训练以提高模型训练的速度。本文将介绍如何在 PyTorch 中进行多 GPU 训练。

安装 PyTorch

首先，确保你的系统上已经安装了 PyTorch。你可以通过以下命令进行安装：

pip install torch torchvision torchaudio

准备数据

在进行多 GPU 训练之前，你需要确保你的数据可以均匀地分布在多个 GPU 上。以下是一个简单的例子：

import torch

# 假设我们有一个数据集，包含 1000 个样本
dataset = torch.randn(1000, 3, 224, 224)

# 将数据集分割成 4 个部分
batch_size = 250
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)

设置多 GPU

接下来，你需要设置你的环境以支持多 GPU。以下是如何使用 torch.nn.DataParallel 来实现多 GPU 训练：

import torch.nn as nn

# 假设我们有一个模型
model = YourModel()

# 检查是否有多个 GPU 可用
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

# 将模型移动到 GPU
model.to('cuda')

训练模型

现在，你可以开始训练模型了。以下是一个简单的训练循环：

optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

for epoch in range(num_epochs):
    for data, target in dataloader:
        # 将数据移动到 GPU
        data, target = data.to('cuda'), target.to('cuda')

        # 前向传播
        output = model(data)

        # 计算损失
        loss = criterion(output, target)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

扩展阅读