PyTorch在Linux上的多GPU训练方法 - 问答

在Linux上使用PyTorch进行多GPU训练，主要依赖于PyTorch的DataParallel或DistributedDataParallel模块。以下是这两种方法的简要介绍和使用步骤：

1. DataParallel

DataParallel是PyTorch中用于多GPU训练的一个简单方法。它会将输入数据分割到多个GPU上，并在每个GPU上并行计算梯度，然后将这些梯度聚合起来更新模型参数。

使用步骤：

确保PyTorch支持多GPU：确保你的PyTorch版本支持多GPU，并且你的系统上安装了多个GPU。
准备数据加载器：使用torch.utils.data.DataLoader来加载数据，并设置num_workers参数以加速数据加载。
将模型移动到GPU：
```
model = YourModel().to('cuda')
```

包装模型：使用DataParallel包装模型：

if torch.cuda.device_count() > 1:
    print(f"Let's use {torch.cuda.device_count()} GPUs!")
    model = nn.DataParallel(model)

训练模型：在训练循环中，像平常一样调用模型的forward方法，并传递输入数据。

2. DistributedDataParallel

DistributedDataParallel是PyTorch中用于多GPU和多节点分布式训练的一个更高级的方法。它提供了更好的性能和可扩展性。

使用步骤：

环境设置：设置环境变量以启用分布式训练：

export MASTER_ADDR='localhost'
export MASTER_PORT='12345'

初始化分布式环境：在代码中初始化分布式环境：

import torch.distributed as dist
dist.init_process_group(backend='nccl')

准备数据加载器：使用torch.utils.data.distributed.DistributedSampler来加载数据，并设置num_replicas和rank参数。

将模型移动到GPU：

model = YourModel().to(torch.device(f'cuda:{rank}'))

包装模型：使用DistributedDataParallel包装模型：

model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])

训练模型：在训练循环中，使用dist.get_rank()来获取当前进程的rank，并根据rank分配数据。

示例代码

以下是一个简单的示例，展示了如何使用DataParallel进行多GPU训练：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 定义模型
class YourModel(nn.Module):
    def __init__(self):
        super(YourModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = torch.relu(torch.max_pool2d(self.conv1(x), 2))
        x = torch.relu(torch.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = torch.relu(self.fc1(x))
        x = torch.dropout(x, training=self.training)
        x = self.fc2(x)
        return torch.log_softmax(x, dim=1)

# 数据加载器
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# 模型和优化器
model = YourModel().to('cuda')
if torch.cuda.device_count() > 1:
    print(f"Let's use {torch.cuda.device_count()} GPUs!")
    model = nn.DataParallel(model)
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(10):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to('cuda'), target.to('cuda')
        optimizer.zero_grad()
        output = model(data)
        loss = nn.functional.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 10 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

通过以上步骤和示例代码，你可以在Linux上使用PyTorch进行多GPU训练。根据你的具体需求和系统配置，选择合适的方法进行训练。

0 赞

0 踩