PyTorch多GPU训练实战：从零实现到ResNet-18模型

本文将介绍如何在PyTorch中实现多GPU训练，涵盖从零开始的手动实现和基于ResNet-18的简洁实现。代码完整可直接运行。

1. 环境准备与库导入

import torch
from torch import nn
from torch.nn import functional as F
from d2l import torch as d2l
from torchvision import models

2. 多GPU参数分发

将模型参数克隆到指定设备并启用梯度计算：

def get_params(params, device):new_params = [p.clone().to(device) for p in params]for p in new_params:p.requires_grad = Truereturn new_params

3. 梯度同步（AllReduce）

实现梯度求和与广播：

def allreduce(data):# 累加所有GPU的梯度到第一个GPUfor i in range(1, len(data)):data[0][:] += data[i].to(data[0].device)# 将结果广播到所有GPUfor i in range(1, len(data)):data[i] = data[0].to(data[i].device)

4. 数据分片

将小批量数据均匀分配到多个GPU：

def split_batch(x, y, devices):assert x.shape[0] == y.shape[0]  # 验证样本数量一致return (nn.parallel.scatter(x, devices),nn.parallel.scatter(y, devices))

5. 训练单个小批量

多GPU训练核心逻辑：

loss = nn.CrossEntropyLoss()def train_batch(x, y, device_params, devices, lr):x_shards, y_shards = split_batch(x, y, devices)  # 数据分片# 计算各GPU损失ls = [loss(net(x_shard, params), y_shard).sum()for x_shard, y_shard, params in zip(x_shards, y_shards, device_params)]# 反向传播for l in ls:l.backward()# 梯度同步with torch.no_grad():for i in range(len(device_params[0])):allreduce([params[i].grad for params in device_params])# 参数更新for param in device_params[0]:d2l.sgd(param, lr, x.shape[0])

6. 完整训练流程

def train(num_gpus, batch_size, lr):train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)devices = [d2l.try_gpu(i) for i in range(num_gpus)]# 初始化模型参数（示例网络）net = nn.Sequential(nn.Conv2d(1, 6, kernel_size=5), nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2),nn.Conv2d(6, 16, kernel_size=5), nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2),nn.Flatten(),nn.Linear(16*4*4, 120), nn.ReLU(),nn.Linear(120, 84), nn.ReLU(),nn.Linear(84, 10))params = list(net.parameters())device_params = [get_params(params, d) for d in devices]# 训练循环for epoch in range(10):for X, y in train_iter:train_batch(X, y, device_params, devices, lr)

7. 简洁实现：修改ResNet-18

def resnet18(num_classes, in_channels=1):def resnet_block(in_channels, out_channels, num_residuals, first_block=False):blk = []for i in range(num_residuals):if i == 0 and not first_block:blk.append(d2l.Residual(in_channels, out_channels, use_1x1conv=False, strides=2))else:blk.append(d2l.Residual(out_channels, out_channels))return nn.Sequential(*blk)# 完整网络结构net = nn.Sequential(nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3),nn.BatchNorm2d(64), nn.ReLU(),nn.MaxPool2d(kernel_size=3, stride=2, padding=1))net.add_module("resnet_block1", resnet_block(64, 64, 2, first_block=True))net.add_module("resnet_block2", resnet_block(64, 128, 2))net.add_module("resnet_block3", resnet_block(128, 256, 2))net.add_module("resnet_block4", resnet_block(256, 512, 2))net.add_module("global_avg_pool", nn.AdaptiveAvgPool2d((1,1)))net.add_module("flatten", nn.Flatten())net.add_module("fc", nn.Linear(512, num_classes))return net# 使用DataParallel包装
net = nn.DataParallel(resnet18(10), device_ids=[0, 1])

8. 运行示例

if __name__ == "__main__":# 从零实现train(num_gpus=2, batch_size=256, lr=0.1)# 简洁实现model = resnet18(10).cuda()model = nn.DataParallel(model, device_ids=[0, 1])