多GPU训练（PyTorch版）

一、多 GPU 训练的三种拆分方式

在多 GPU 训练中，有三种拆分训练任务的方式，各有优缺点：

1. 网络拆分（层间拆分）：把网络的不同层分配到不同 GPU，每个 GPU 处理一部分层的计算。优点是可以用更大的网络，每个 GPU 的显存占用小；缺点是 GPU 之间需要密集同步，层间数据传输量大，同步难度高，不推荐。

2. 层内拆分：把一层的计算任务拆分到多个 GPU，比如卷积层的通道拆分，全连接层的输出单元拆分。优点是可以处理更大的网络，显存线性扩展；缺点是需要大量同步操作，数据传输量更大，也不推荐。

3. 数据并行（数据拆分）：把小批量数据分成 k 份（k 是 GPU 数量），每个 GPU 拿到一份数据，独立计算损失和梯度，然后聚合所有 GPU 的梯度，再把聚合后的梯度分发到每个 GPU 更新参数。优点是最简单，容易实现，同步只需要在每个小批量处理后进行，是目前最推荐的方式，只要 GPU 显存足够就用这种。

二、数据并行的训练流程

假设有 k 个 GPU，数据并行的训练步骤是：

1. 每次训练迭代，把小批量样本均匀分成 k 份，分配到每个 GPU 上。

2. 每个 GPU 用自己的那部分数据，计算模型的损失和参数的梯度。

3. 把 k 个 GPU 上的局部梯度聚合起来，得到这个小批量的随机梯度。

4. 把聚合后的梯度分发到每个 GPU。

5. 每个 GPU 用这个聚合后的梯度，更新自己维护的完整模型参数。

注意：用 k 个 GPU 训练时，小批量的大小要设置成原来的 k 倍，这样每个 GPU 的工作量和单 GPU 训练时一样；同时可能需要提高学习率，还有批量归一化的参数要每个 GPU 单独维护。

三、PyTorch 实现多 GPU 训练的步骤

1. 定义简单模型（修改版 LeNet）

首先定义一个简单的卷积神经网络，和 LeNet 类似，用于演示多 GPU 训练：

import torch
from torch import nn
import d2l

# 定义LeNet模型
def lenet(X, params):
    # 卷积层1
    h1_conv = nn.functional.conv2d(input=X, weight=params[0], bias=params[1])
    h1_activation = nn.functional.relu(h1_conv)
    h1 = nn.functional.avg_pool2d(input=h1_activation, kernel_size=2, stride=2)
    # 卷积层2
    h2_conv = nn.functional.conv2d(input=h1, weight=params[2], bias=params[3])
    h2_activation = nn.functional.relu(h2_conv)
    h2 = nn.functional.avg_pool2d(input=h2_activation, kernel_size=2, stride=2)
    # 展平
    h2 = h2.reshape(h2.shape[0], -1)
    # 全连接层1
    h3_linear = nn.functional.linear(input=h2, weight=params[4], bias=params[5])
    h3 = nn.functional.relu(h3_linear)
    # 全连接层2
    y_hat = nn.functional.linear(input=h3, weight=params[6], bias=params[7])
    return y_hat

# 初始化模型参数
scale = 0.01
W1 = torch.randn(size=(20, 1, 3, 3)) * scale
b1 = torch.zeros(20)
W2 = torch.randn(size=(50, 20, 5, 5)) * scale
b2 = torch.zeros(50)
W3 = torch.randn(size=(800, 128)) * scale
b3 = torch.zeros(128)
W4 = torch.randn(size=(128, 10)) * scale
b4 = torch.zeros(10)
params = [W1, b1, W2, b2, W3, b3, W4, b4]

# 交叉熵损失函数
loss = nn.CrossEntropyLoss(reduction='none')

2. 数据同步操作

多 GPU 训练需要两个关键操作：把参数复制到各个 GPU，以及聚合各个 GPU 的梯度。

（1）复制参数到 GPU

定义get_params函数，把模型参数复制到指定的 GPU，并开启梯度跟踪：

def get_params(params, device):
    new_params = [p.to(device) for p in params]
    for p in new_params:
        p.requires_grad_()
    return new_params

（2）聚合梯度（allreduce）

定义allreduce函数，把各个 GPU 上的梯度相加，然后把结果广播到所有 GPU：

def allreduce(data):
    # 把所有GPU上的数据加到第一个GPU上
    for i in range(1, len(data)):
        data[0][:] += data[i].to(data[0].device)
    # 把聚合后的结果复制到其他GPU
    for i in range(1, len(data)):
        data[i][:] = data[0].to(data[i].device)

3. 数据分发

定义split_batch函数，把数据和标签均匀拆分到各个 GPU 上：

def split_batch(X, y, devices):
    """将X和y拆分到多个设备上"""
    assert X.shape[0] == y.shape[0]
    return (nn.parallel.scatter(X, devices),
            nn.parallel.scatter(y, devices))

4. 单批次训练

定义train_batch函数，处理一个小批量的多 GPU 训练：

def train_batch(X, y, device_params, devices, lr):
    # 把数据拆分到各个GPU
    X_shards, y_shards = split_batch(X, y, devices)
    # 在每个GPU上分别计算损失
    ls = [loss(lenet(X_shard, device_W), y_shard).sum()
          for X_shard, y_shard, device_W in zip(
              X_shards, y_shards, device_params)]
    # 在每个GPU上分别执行反向传播
    for l in ls:
        l.backward()
    # 聚合所有GPU的梯度
    with torch.no_grad():
        for i in range(len(device_params[0])):
            allreduce(
                [device_params[c][i].grad for c in range(len(devices))])
    # 在每个GPU上分别更新模型参数
    for param in device_params:
        d2l.sgd(param, lr, X.shape[0])  # 使用全小批量的大小来更新

5. 完整训练流程

定义train函数，完成整个训练过程：

def train(num_gpus, batch_size, lr):
    # 加载Fashion-MNIST数据集
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    # 获取所有GPU设备
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    # 把模型参数复制到每个GPU
    device_params = [get_params(params, d) for d in devices]
    num_epochs = 10
    # 用于可视化训练过程
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    timer = d2l.Timer()
    for epoch in range(num_epochs):
        timer.start()
        # 训练每个小批量
        for X, y in train_iter:
            train_batch(X, y, device_params, devices, lr)
        # 等待所有GPU计算完成
        torch.cuda.synchronize()
        timer.stop()
        # 在GPU0上评估模型精度
        animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(
            lambda x: lenet(x, device_params[0]), test_iter, devices[0]),))
    print(f'测试精度：{animator.Y[0][-1]:.2f}，{timer.avg():.1f}秒/轮，'
          f'在{str(devices)}')

6. 训练示例

（1）单 GPU 训练

train(num_gpus=1, batch_size=256, lr=0.2)

运行结果大概是：测试精度 0.84，2.7 秒 / 轮，在[device(type='cuda', index=0)]

（2）双 GPU 训练

train(num_gpus=2, batch_size=256, lr=0.2)

运行结果大概是：测试精度 0.83，3.6 秒 / 轮，在[device(type='cuda', index=0), device(type='cuda', index=1)]

四、小结

1. 多 GPU 训练中，数据并行是最简单实用的方式，只要 GPU 显存足够就优先选择。

2. 数据并行的核心是把数据拆分到各个 GPU，独立计算梯度后聚合，再更新参数。

3. 使用 k 个 GPU 时，小批量大小要设置为原来的 k 倍，同时可能需要调整学习率。

4. PyTorch 中可以通过nn.parallel.scatter拆分数据，通过自定义allreduce函数聚合梯度，实现多 GPU 训练。

（注：文档部分内容可能由 AI 生成）