简单分布式模型训练（多GPU）

要在模型训练中使用两个 GPU，你可以利用 PyTorch 的 DataParallel 或 torch.nn.parallel.DistributedDataParallel 来实现多 GPU 并行计算。以下是如何在 PyTorch 中使用两个 GPU 的示例：

1. 使用 torch.nn.DataParallel

DataParallel 是 PyTorch 提供的一个较为简单的多 GPU 并行训练方法。它可以将模型复制到多个 GPU 上，并将数据批次分发给每个 GPU，然后将各个 GPU 的结果合并。

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset

# 示例：创建一个简单的数据集
class SimpleDataset(Dataset):
    def __init__(self, size):
        self.size = size
        self.data = torch.randn(size, 10)
        self.labels = torch.randint(0, 2, (size,))

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# 简单的神经网络模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 2)

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

# 数据集和DataLoader
dataset = SimpleDataset(1000)
data_loader = DataLoader(dataset, batch_size=32, num_workers=2)

# 创建模型、损失函数和优化器
model = SimpleModel()

# 使用DataParallel并行化模型，指定GPU设备（device_ids=[0, 1]表示使用两个GPU）
model = nn.DataParallel(model, device_ids=[0, 1])  # 使用GPU 0和GPU 1
model = model.to('cuda')  # 将模型移动到GPU

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

# 训练循环
for epoch in range(10):
    for data, labels in data_loader:
        # 将数据移动到GPU
        data, labels = data.to('cuda'), labels.to('cuda')
        
        # 正向传播
        outputs = model(data)
        loss = criterion(outputs, labels)
        
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch+1}, Loss: {loss.item()}")

解释：

1. nn.DataParallel(model, device_ids=[0, 1])：将模型并行化到指定的多个 GPU 上（在此例中为 GPU 0 和 GPU 1）。
2. model.to('cuda')：将模型移动到 GPU 上（如果你的机器有多个 GPU，DataParallel 会自动将模型复制到多个 GPU 上）。
3. data.to('cuda') 和 labels.to('cuda')：将输入数据和标签移动到 GPU 上进行计算。

2. 使用 torch.nn.parallel.DistributedDataParallel

DistributedDataParallel 是一种更加高效的多 GPU 并行方法，特别适用于大规模训练。在分布式环境下使用时，它能显著提升训练速度和效率。以下是使用 DistributedDataParallel 的示例：

步骤：

1. 每个进程在不同的 GPU 上启动。
2. 使用 torch.distributed.launch 或 torch.multiprocessing.spawn 来启动多个进程。

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import os

# 初始化分布式环境
def init_process(rank, size, fn, backend='nccl'):
    """ 初始化分布式环境 """
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group(backend, rank=rank, world_size=size)
    fn(rank, size)

# 示例：创建一个简单的数据集
class SimpleDataset(Dataset):
    def __init__(self, size):
        self.size = size
        self.data = torch.randn(size, 10)
        self.labels = torch.randint(0, 2, (size,))

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# 简单的神经网络模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 2)

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

# 模型训练函数
def train(rank, size):
    # 数据集和DataLoader
    dataset = SimpleDataset(1000)
    data_loader = DataLoader(dataset, batch_size=32, num_workers=2, shuffle=True)

    # 创建模型、损失函数和优化器
    model = SimpleModel().to(rank)
    model = DDP(model, device_ids=[rank])  # 将模型移动到指定GPU并使用DDP进行分布式训练

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

    # 训练循环
    for epoch in range(10):
        for data, labels in data_loader:
            data, labels = data.to(rank), labels.to(rank)

            # 正向传播
            outputs = model(data)
            loss = criterion(outputs, labels)

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

        print(f"Rank {rank}, Epoch {epoch+1}, Loss: {loss.item()}")

# 使用torch.multiprocessing.spawn启动多进程
def main():
    size = 2  # 使用两个GPU
    torch.multiprocessing.spawn(init_process, args=(size, train), nprocs=size, join=True)

if __name__ == "__main__":
    main()

解释：

1. DistributedDataParallel：将模型分布式并行到多个 GPU 上。通过 device_ids=[rank] 将每个进程绑定到一个特定的 GPU 上（rank 是每个进程的 GPU 索引）。
2. dist.init_process_group：初始化分布式训练环境。
3. torch.multiprocessing.spawn：启动多个进程，每个进程负责训练一个模型副本并与其他进程进行梯度同步。

总结：

• DataParallel：适用于较小规模的多 GPU 环境，容易实现并且代码较为简单。
• DistributedDataParallel：适用于大规模的多 GPU 环境，能够提供更好的性能和扩展性，尤其是在多节点训练中更为高效。

你可以根据硬件环境的规模和训练的复杂度来选择使用 DataParallel 或 DistributedDataParallel。