文档首页/ AI开发平台ModelArts/ ModelArts用户指南（Standard）/ 使用ModelArts Standard训练模型/ 分布式模型训练/ 创建单机多卡的分布式训练（DataParallel）
更新时间：2024-08-16 GMT+08:00
查看PDF
分享

创建单机多卡的分布式训练（DataParallel）

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

MindSpore引擎的分布式训练参见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 ###
父主题： 分布式模型训练

相关文档