34 多GPU训练实现 [动手学深度学习v2]-优快云博客

本文链接：https://blog.youkuaiyun.com/gpx33333/article/details/122010385

本文介绍了一种使用多个GPU进行并行训练的方法，包括参数复制、梯度聚合等关键步骤，并展示了如何利用PyTorch实现多GPU训练。同时，文中还讨论了不同设置下训练效果的影响因素。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

多GPU训练实现

import torch
from torch import nn

# 将参数复制到其他GPU上
def get_params(params, device):
    new_params = [p.clone().to(device) for p in params]
    for p in new_params:
        p.requires_grad_()
    return new_params

# 将所有向量相加，并将结果广播给所有GPU
def allreduce(data):
    for i in range(1, len(data)):
        data[0][:] += data[i].to(data[0].device)
    for i in range(1, len(data)):
        data[i] = data[0].to(data[i].device)
        
# 将一个小批量数据均匀地分布在多个GPU上面
data = torch.arange(20).reshape(4, 5)
devices = [torch.device('cuda:0'), torch.device('cuda:1')]
split = nn.parallel.scatter(data, devices)

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))

# 在一个小批量上实现多GPU训练
def train_batch(X, y, device_params, devices, lr):
    X_shards, y_shards = split(X, y, devices)
    ls = [loss(model(X_shard, device_W), y_shard).sum() for X_shard, y_shard, device_W in zip(X_shards, y_shards, device_params)]
    for l in ls:
        l.backword()
    with torch.no_grad():
        for i in range(len(devices_params[0])):
            allreduce([device_params[c][i].grad for c in range(len(devices))])
    for param in device_params:
        SGD(param, lr, X.shape[0])

多GPU的简洁实现

import torch
from torch import nn

# 相比于单GPU训练，只用修改这一行就行
net = nn.DataParallel(net, device_ids=devices)

对于精度来说，batch_size=1通常是最好的情况【QA】。
验证集准确率震荡较大是lr影响最大【QA】。
batch_size增大，lr也要增大【QA】。
如果有网络中batch normlization，lr可以稍微调大一点【QA】。