多GPU训练数据并行k个GPU的训练过程以我们之前定义的LeNet模型为例首先我们需要向多个设备分发参数并附加梯度get_params。 如果没有参数就不可能在GPU上评估网络。 第二需要跨多个设备对参数求和也就是说需要一个allreduce函数。两个关键函数的实现*每个 GPU 都需要有一份完整的模型参数副本将所有 GPU 上的向量通常是梯度相加求和然后把结果广播回所有 GPU先复制到同一个GPU上相加然后将结果复制到其余的GPU上将一个小批量数据均匀地分布在多个GPU上使用nn.parallel.scatter函数带标签的数据训练训练函数1拆分数据将 X, y 均匀分布到各 GPU2各 GPU 独立前向传播 计算损失3各 GPU 独立反向传播4梯度聚合所有 GPU 的梯度相加广播回所有 GPU5各 GPU 独立更新参数SGD改进方向多GPU可以并行这里写法看似顺序整体训练模型load加载数据获取GPU列表将模型参数复制到所有GPU为每个小批量执行多GPU训练简洁代码实现nn.DataParallel