PyTorch在Linux上的多GPU訓練方法

在Linux上使用PyTorch進行多GPU訓練，主要依賴于PyTorch的DataParallel或DistributedDataParallel模塊。以下是這兩種方法的簡要介紹和使用步驟：

1. DataParallel

DataParallel是PyTorch中用于多GPU訓練的一個簡單方法。它會將輸入數據分割到多個GPU上，并在每個GPU上并行計算梯度，然后將這些梯度聚合起來更新模型參數。

使用步驟：

1. 確保PyTorch支持多GPU： 確保你的PyTorch版本支持多GPU，并且你的系統上安裝了多個GPU。

2. 準備數據加載器： 使用torch.utils.data.DataLoader來加載數據，并設置num_workers參數以加速數據加載。

3. 將模型移動到GPU：

   model = YourModel().to('cuda')
   

4. 包裝模型： 使用DataParallel包裝模型：

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

5. 訓練模型： 在訓練循環中，像平常一樣調用模型的forward方法，并傳遞輸入數據。

2. DistributedDataParallel

DistributedDataParallel是PyTorch中用于多GPU和多節點分布式訓練的一個更高級的方法。它提供了更好的性能和可擴展性。

使用步驟：

1. 環境設置： 設置環境變量以啟用分布式訓練：

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

2. 初始化分布式環境： 在代碼中初始化分布式環境：

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

3. 準備數據加載器： 使用torch.utils.data.distributed.DistributedSampler來加載數據，并設置num_replicas和rank參數。

4. 將模型移動到GPU：

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

5. 包裝模型： 使用DistributedDataParallel包裝模型：

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

6. 訓練模型： 在訓練循環中，使用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訓練。根據你的具體需求和系統配置，選擇合適的方法進行訓練。