此文章是对李沐计算性能章节的理解
多卡计算的方案
主要有:
- 网络分区
- 按层分区
- 数据并行
数据并行比较常见,目前多卡训练都是这种数据并行的方式实现的。
数据并行
数据并行是指的将一个完整批次的数据分成几个小批次分到几块GPU运算可以提升显存大小,提升训练批次大小,减小训练时间(当批次很大时,一块卡运行载荷爆满会导致训练速度变慢)。
主要要解决的问题是不同卡之间梯度如何保持一致。
步骤:
- 将批次的数据平均分到多卡上
- 将模型完整的复制到多卡上
- 在每个卡上计算局部损失和局部梯度
- 将每个卡计算的梯度都汇集到一起并再返回到每个卡上(相当于每个卡都计算了了一整个批次大小数据的梯度)
- 每个卡单独更新参数(参数都一样因为是一个模型复制多份;全局梯度一样,所以最终更新后的每个卡中模型的参数都是一样的所有卡中的模型均是同一个)
代码实现
- 单机多卡:
一行代码搞定,只需要对net进行一下以下包装就可以。
net = nn.DataParallel(net, device_ids=devices)
- 多机多卡(分布式?)
埋个坑