详细参考《动手学深度学习》P233,8.4节多GPU计算。
数据并行的方法
- 把一个batch的所有数据平均分配到每块GPU的显存里,把模型参数在每个GPU显存上拷贝一份。
- 每个GPU根据自己所分到的数据,计算本地梯度
- 所有GPU的本地梯度相加(这里的梯度只是模型参数的梯度,不包括中间变量的梯度),得到总的梯度。(注意,梯度不用求平均,因为更新模型参数的时候,会除以batch_size,在那里做了平均)
- 把总的梯度广播到每块GPU上,每块gpu用总的梯度更新模型参数。
for i in len(gpu_params[0])):
allreduce((gpu[params[c][i].grad for c in range(len(ctx))])
如果把batch_size调大1倍,则学习率也需要调大一倍,因为模型更新的次数少了一倍。
CPU GPU同步数据
注意显存和内存的数据是存储在不同的地方,必须通过某些内置函数才能实现两者之间的数据拷贝。
比如训练样本数据,用CPU处理后,就需要把它拷贝到GPU上。
各个GPU也只能访问自己显存内的数据。