文档首页/
AI开发平台ModelArts/
ModelArts用户指南(Standard)/
使用ModelArts Standard训练模型/
分布式模型训练/
创建单机多卡的分布式训练(DataParallel)
更新时间:2024-10-24 GMT+08:00
创建单机多卡的分布式训练(DataParallel)
本章节介绍基于PyTorch引擎的单机多卡数据并行训练。
MindSpore引擎的分布式训练参见MindSpore官网。
训练流程简述
单机多卡数据并行训练流程介绍如下:
- 将模型复制到多个GPU上
- 将一个Batch的数据均分到每一个GPU上
- 各GPU上的模型进行前向传播,得到输出
- 主GPU(逻辑序号为0)收集各GPU的输出,汇总后计算损失
- 分发损失,各GPU各自反向传播梯度
- 主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 ###
父主题: 分布式模型训练