数据并行：torch.nn.DataParallel的数据并行原理

最新推荐文章于 2024-06-10 00:44:25 发布

飞翔の荷兰人

最新推荐文章于 2024-06-10 00:44:25 发布

阅读量709

点赞数

分类专栏：分布式训练文章标签：深度学习 pytorch 机器学习分布式

本文链接：https://blog.csdn.net/Corleone26/article/details/130069433

版权

分布式训练专栏收录该内容

1 篇文章 0 订阅

订阅专栏

文章探讨了PyTorch中使用DataParallel进行多卡训练的机制，强调了在单卡和多卡训练中调用backward()和optimizer.step()的一致性。优化器通过all-reduce算法聚合梯度并在主GPU上更新参数，确保多GPU间模型参数同步。此外，提到了自定义数据并行实现的可能性，但PyTorch自带的方案在效率和易用性之间找到了平衡。

摘要由CSDN通过智能技术生成

直接使用沐神d2l的代码作为示例，可以看到多卡数据并行的代码与直接单卡训练几乎没有改动

def train(net, 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)]
    def init_weights(m):
        if type(m) in [nn.Linear, nn.Conv2d]:
            nn.init.normal_(m.weight, std=0.01)
    net.apply(init_weights)
    # 在多个GPU上设置模型
    net = nn.DataParallel(net, device_ids=devices)
    trainer = torch.optim.SGD(net.parameters(), lr)
    loss = nn.CrossEntropyLoss()
    timer, num_epochs = d2l.Timer(), 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    for epoch in range(num_epochs):
        net.train()
        timer.start()
        for X, y in train_iter:
            trainer.zero_grad()
            X, y = X.to(devices[0]), y.to(devices[0]) # nn.DataParallel指定需要先把参数放到devices[0]
            l = loss(net(X), y) 
            l.backward() 
            trainer.step() 
        timer.stop()
        animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(net, test_iter),))
    print(f'测试精度：{animator.Y[0][-1]:.2f}，{timer.avg():.1f}秒/轮，'
          f'在{str(devices)}')

从用户的角度来看，在单卡和多卡训练中调用backward()和optimizer.step()的方式是完全一样的
而PyTorch底层会自动处理多GPU之间的数据传输和参数更新等细节。
在这个代码中，虽然output和loss是在GPU0上计算的
但是backward会根据计算图，在不同卡的参数上更新属于自己那部分的梯度
而梯度汇聚和参数更新，都是由trainer.step()这一步操作完成的。
在调用backward()计算梯度之后，trainer.step()会执行all-reduce操作，
将各个GPU上计算得到的梯度加和，并在主GPU上更新模型参数，然后将更新后的参数分发到每个GPU上。
当然数据并行也可以选择主GPU分发梯度(直接接受梯度，加和，分发，标准的all_reduce)，这样每个GPU分别更新参数，理论上效果相同。但这种分发参数的方式，chatgpt认为更加同步
因此，trainer.step()既包括了梯度汇聚也包括了参数更新。

PyTorch中的优化器是负责更新模型参数的。在DataParallel中，每个GPU上的模型参数都是完整的，但是每个GPU只计算了一部分数据的梯度。
优化器会收集所有GPU上的梯度，然后使用all-reduce算法将它们加和，并在主GPU上更新模型参数。
更新后的参数会通过分发机制发送到所有GPU上，确保各个GPU上的参数始终保持一致。
这段在代码中，这个过程被封装在了trainer.step()这一语句中。

其实，也有其他的数据并行实现方法，比如沐神在前一章节的代码中自己实现的方法：
把输入和label都分别分发到不同的卡上，然后每个卡可以分别计算自己的loss，然后all_reduce一次loss(简单的加和再分发)，之后每张卡分别更新自己的梯度即可。
torch.nn.DataParallel则是把每张卡的输出聚合到GPU0上，然后在GPU0上与label计算loss，根据计算图反向传播，让每张卡上获得自己的梯度。优化器则对梯度进行聚合，在主GPU更新模型参数，再把新的参数分发到每个GPU。
直观感觉明显是第一种效率更高，而且实现简单。第二种对优化器和计算图都有要求，关键是反向传播是从主GPU上往不同卡上反向传误差啊！！！但是对用户友好，几乎不用改代码。关键是，实验起来，也是使用pytorch自己的这种方式效率高，底层应该有很多优化。
写到这里，真诚的希望大佬们给一些指点。