pytorch train多GPU訓練

在PyTorch中，使用多GPU進行訓練可以顯著提高訓練速度。以下是一個簡單的示例，展示了如何使用torch.nn.DataParallel或torch.nn.parallel.DistributedDataParallel進行多GPU訓練。

使用 torch.nn.DataParallel

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 定義一個簡單的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(784, 10)

    def forward(self, x):
        return self.fc(x)

# 創建模型實例
model = SimpleModel()

# 使用 DataParallel 包裝模型
if torch.cuda.device_count() > 1:
    print("Using", torch.cuda.device_count(), "GPUs")
    model = nn.DataParallel(model)

# 將模型放到GPU上
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

# 定義損失函數和優化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 數據預處理
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

# 加載數據集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

# 訓練模型
num_epochs = 10
for epoch in range(num_epochs):
    for data, target in train_loader:
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")

使用 torch.nn.parallel.DistributedDataParallel

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import torch.distributed as dist
import torch.multiprocessing as mp

# 定義一個簡單的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(784, 10)

    def forward(self, x):
        return self.fc(x)

def train(rank, world_size):
    # 初始化進程組
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

    # 創建模型實例
    model = SimpleModel()

    # 使用 DistributedDataParallel 包裝模型
    model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])

    # 將模型放到GPU上
    device = torch.device(f"cuda:{rank}")
    model.to(device)

    # 定義損失函數和優化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)

    # 數據預處理
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

    # 加載數據集
    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, sampler=train_sampler)

    # 訓練模型
    num_epochs = 10
    for epoch in range(num_epochs):
        train_sampler.set_epoch(epoch)
        for data, target in train_loader:
            data, target = data.to(device), target.to(device)

            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

        print(f"Rank {rank}, Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")

    # 結束進程組
    dist.destroy_process_group()

def main():
    world_size = torch.cuda.device_count()
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)

if __name__ == "__main__":
    main()

說明

1. torch.nn.DataParallel：適用于單機多GPU環境，簡單且易于使用。
2. torch.nn.parallel.DistributedDataParallel：適用于分布式環境，性能更優，但配置和使用相對復雜。

選擇哪種方法取決于你的具體需求和環境。