gpu并行训练

发表于 | 分类于

nn.DataParallel

若只考虑单机多卡并行和数据平行训练时，最简单的方法是用 nn.DataParallel。

先设置所使用的GPU卡，然后用 nn.DataParallel 包装模型。

1
2
3
4
import os
os.environ["CUDA_VISIBLE_DEVICES"] ="0, 2"# 使用0和2号卡
model=Classifier()
model =nn.DataParallel(model)

Pytorch的nn.DataParallel - 知乎 (zhihu.com)

前向过程：

  1. 分发mini-batch到每个GPU上，
  2. 将模型复制到每一个GPU上，
  3. 模型对mini-batch进行forward计算结果，
  4. 将每个GPU上的计算结果汇总到第一个GPU上。

反向传播过程：

  1. 在第一个GPU上根据计算的loss计算梯度，
  2. 将第一个GPU上计算的梯度值分发到每个GPU上，
  3. 每个GPU上进行梯度更新，
  4. 将每个GPU上更新的梯度值汇总到第一个GPU上。
    其中foward过程都会把第一个GPU上的模型重新分发到每个每个GPU上。各个卡只计算到loss，然后0号卡做loss平均，最后分发到各个卡上求梯度并进行参数更新。

文章有错误：应该是
1：把每张卡的输出gather到GPU0，
2：只在GPU0在算loss，
3：把loss scatter到其他GPU上，
4：每个GPU算各自的梯度，再把梯度汇总到GPU0上

关于这个流程不同的文章有不同的说法。

nn.parallel.DistributedDataParallel

PyTorch 源码解读之分布式训练了解一下？ - 知乎 (zhihu.com)

image-20220112162455330

pytorch多GPU并行训练教程_哔哩哔哩_bilibili

image-20220112162623305

image-20220112163216347

image-20220112163334093

同步bn

nn.DataParallel

单机多卡。单进程多线程。工作与单机。

nn.parallel.DistributedDataParallel

可以多级多卡，速度快于DataParallel。

image-20220112170831371

image-20220112171028243

多gpu训练，每块gpu上初始模型需要一摸一样。

loss_reduce,多块gpu loss做平均。