多GPU训练:数据并行

多GPU训练:数据并行

多GPU训练
数据并行,k个GPU的训练过程:

以我们之前定义的LeNet模型为例:

首先,我们需要向多个设备分发参数并附加梯度(get_params)。 如果没有参数,就不可能在GPU上评估网络。 第二,需要跨多个设备对参数求和,也就是说,需要一个allreduce函数。

两个关键函数的实现:

*每个 GPU 都需要有一份完整的模型参数副本

将所有 GPU 上的向量(通常是梯度)相加求和,然后把结果广播回所有 GPU

先复制到同一个GPU上相加,然后将结果复制到其余的GPU上

将一个小批量数据均匀地分布在多个GPU上,使用nn.parallel.scatter函数

带标签的数据

训练:
训练函数


(1)拆分数据:将 X, y 均匀分布到各 GPU

(2)各 GPU 独立前向传播 + 计算损失

(3)各 GPU 独立反向传播

(4)梯度聚合:所有 GPU 的梯度相加,广播回所有 GPU

(5)各 GPU 独立更新参数(SGD)
改进方向:多GPU可以并行(这里写法看似顺序)

整体训练模型:
load加载数据


获取GPU列表


将模型参数复制到所有GPU


为每个小批量执行多GPU训练



简洁代码实现nn.DataParallel