更新时间:2025-08-07 GMT+08:00

创建单机多卡的分布式训练(DataParallel)

在深度学习领域,随着模型规模的不断扩大,训练时间也随之增加。为了提高训练效率,需要采用高效的并行计算方法。在单机环境下,如何充分利用多块GPU卡的计算能力成为一个关键问题。本章节将介绍基于PyTorch引擎的单机多卡数据并行训练方法,通过合理的数据划分和模型同步策略,充分发挥多GPU的计算能力,显著提升训练效率。

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 ###