更新时间:2024-08-16 gmt 08:00

创建单机多卡的分布式训练(dataparallel)-九游平台

本章节介绍基于pytorch引擎的单机多卡数据并行训练。

mindspore引擎的分布式训练参见。

训练流程简述

单机多卡数据并行训练流程介绍如下:

  1. 将模型复制到多个gpu上
  2. 将一个batch的数据均分到每一个gpu上
  3. 各gpu上的模型进行前向传播,得到输出
  4. 主gpu(逻辑序号为0)收集各gpu的输出,汇总后计算损失
  5. 分发损失,各gpu各自反向传播梯度
  6. 主gpu收集梯度并更新参数,将更新后的模型参数分发到各gpu

具体流程图如下:

图1 单机多卡数据并行训练

代码改造点

模型分发:dataparallel(model)

完整代码由于代码变动较少,此处进行简略介绍。

import torch
class net(torch.nn.module):
	pass
model = net().cuda()
### dataparallel begin ###
model = torch.nn.dataparallel(net().cuda())
### dataparallel end ###

相关文档

网站地图