详解如何使用Pytorch进行多卡训练

目录
  • 1.DP
  • 2.DDP
    • 2.1Pytorch分布式基础
    • 2.2Pytorch分布式训练DEMO

当一块GPU不够用时,我们就需要使用多卡进行并行训练。其中多卡并行可分为数据并行和模型并行。具体区别如下图所示:

由于模型并行比较少用,这里只对数据并行进行记录。对于pytorch,有两种方式可以进行数据并行:数据并行(DataParallel, DP)和分布式数据并行(DistributedDataParallel, DDP)。

在多卡训练的实现上,DP与DDP的思路是相似的:

1、每张卡都复制一个有相同参数的模型副本。

2、每次迭代,每张卡分别输入不同批次数据,分别计算梯度。

3、DP与DDP的主要不同在于接下来的多卡通信:

DP的多卡交互实现在一个进程之中,它将一张卡视为主卡,维护单独模型优化器。所有卡计算完梯度后,主卡汇聚其它卡的梯度进行平均并用优化器更新模型参数,再将模型参数更新至其它卡上。

DDP则分别为每张卡创建一个进程,每个进程相应的卡上都独立维护模型和优化器。在每次每张卡计算完梯度之后,进程之间以NCLL(NVIDIA GPU通信)为通信后端,使各卡获取其它卡的梯度。各卡对获取的梯度进行平均,然后执行后续的参数更新。由于每张卡上的模型与优化器参数在初始化时就保持一致,而每次迭代的平均梯度也保持一致,那么即使没有进行参数复制,所有卡的模型参数也是保持一致的。

Pytorch官方推荐我们使用DDP。DP经过我的实验,两块GPU甚至比一块还慢。当然不同模型可能有不同的结果。下面分别对DP和DDP进行记录。

1.DP

Pytorch的DP实现多GPU训练十分简单,只需在单GPU的基础上加一行代码即可。以下是一个DEMO的代码。

import torch
from torch import nn
from torch.optim import Adam
from torch.nn.parallel import DataParallel

class DEMO_model(nn.Module):
  def __init__(self, in_size, out_size):
    super().__init__()
    self.fc = nn.Linear(in_size, out_size)
  def forward(self, inp):
    outp = self.fc(inp)
    print(inp.shape, outp.device)
    return outp
model = DEMO_model(10, 5).to('cuda')
model = DataParallel(model, device_ids=[0, 1]) # 额外加这一行
adam = Adam(model.parameters())
# 进行训练
for i in range(1):
  x = torch.rand([128, 10]) # 获取训练数据,无需指定设备
  y = model(x) # 自动均匀划分数据批量并分配至各GPU,输出结果y会聚集到GPU0中
  loss = torch.norm(y)
  loss.backward()
  adam.step()

其中model = DataParallel(model, device_ids=[0, 1])这行将模型复制到0,1号GPU上。输入数据x无需指定设备,它将会被均匀分配至各块GPU模型,进行前向传播。之后各块GPU的输出再合并到GPU0中,得到输出y。输出y在GPU0中计算损失,并进行反向传播计算梯度、优化器更新参数。

2.DDP

为了对分布式编程有基本概念,首先使用pytorch内部的方法实现一个多进程程序,再使用DDP模块实现模型的分布式训练。

2.1Pytorch分布式基础

首先使用pytorch内部的方法编写一个多进程程序作为编写分布式训练的基础。

import os, torch
import torch.multiprocessing as mp
import torch.distributed as dist
def run(rank, size):
  tensor = torch.tensor([1,2,3,4], device='cuda:'+str(rank)) # ——1——
  group = dist.new_group(range(size)) # ——2——
  dist.all_reduce(tensor=tensor, group=group, op=dist.ReduceOp.SUM) # ——3——
  print(str(rank)+ ': ' + str(tensor) + '\n')
def ini_process(rank, size, fn, backend = 'nccl'):
  os.environ['MASTER_ADDR'] = '127.0.0.1' # ——4——
  os.environ['MASTER_PORT'] = '1234'
  dist.init_process_group(backend, rank=rank, world_size=size) # ——5——
  fn(rank, size) # ——6——
if __name__ == '__main__': # ——7——
  mp.set_start_method('spawn') # ——8——
  size = 2 # ——9——
  ps = []
  for rank in range(size):
    p = mp.Process(target=ini_process, args=(rank, size, run)) # ——10——
    p.start()
    ps.append(p)
  for p in ps: # ——11——
    p.join()

以上代码主进程创建了两个子进程,子进程之间使用NCCL后端进行通信。每个子进程各占用一个GPU资源,实现了所有GPU张量求和的功能。细节注释如下:

1、为每个子进程定义相同名称的张量,并分别分配至不同的GPU,从而能进行后续的GPU间通信。

2、定义一个通信组,用于后面的all_reduce通信操作。

3、all_reduce操作以及其它通信方式请看下图:

4、定义编号(rank)为0的ip和端口地址,让每个子进程都知道。ip和端口地址可以随意定义,不冲突即可。如果不设置,子进程在涉及进程通信时会出错。

5、初始化子进程组,定义进程间的通信后端(还有GLOO、MPI,只有NCCL支持GPU间通信)、子进程rank、子进程数量。只有当该函数在size个进程中被调用时,各进程才会继续从这里执行下去。这个函数统一了各子进程后续代码的开始时间。

6、执行子进程代码。

7、由于创建子进程会执行本程序,因此主进程的执行需要放在__main__里,防止子进程执行。

8、开始创建子进程的方式:spawn、fork。windows默认spawn,linux默认fork。具体区别请百度。

9、由于是以NCCL为通信后端的分布式训练,如果不同进程中相同名称的张量在同一GPU上,当这个张量进行进程间通信时就会出错。为了防止出错,限制每张卡独占一个进程,每个进程独占一张卡。这里有两张卡,所以最多只能创建两个进程。

10、创建子进程,传入子进程的初始化方法,及子进程调用该方法的参数。

11、等待子进程全部运行完毕后再退出主进程。

输出结果如下:

正是各进程保存在不同GPU上的张量的广播求和(all_reduce)的结果。

参考:https://pytorch.org/tutorials/intermediate/dist_tuto.html

2.2Pytorch分布式训练DEMO

我们实际上可以根据上面的分布式基础写一个分布式训练,但由于不知道pytorch如何实现GPU间模型梯度的求和,即官方教程中所谓的ring_reduce(没找到相关API),时间原因,就不再去搜索相关方法了。这里仅记录pytorh内部的分布式模型训练,即利用DDP模块实现。Pytorch版本1.12.1。

import torch,os
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
from torch import nn
def example(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)  # ——1——
    model = nn.Linear(2, 1, False).to(rank)
    if rank == 0: # ——2——
        model.load_state_dict(torch.load('model_weight'))
    # model_stat = torch.load('model_weight', {'cuda:0':'cuda:%d'%rank})  #这样读取保险一点
    # model.load_state_dict(model_stat)
    opt = optim.Adam(model.parameters(), lr=0.0001) # ——3——
    opt_stat = torch.load('opt_weight', {'cuda:0':'cuda:%d'%rank}) # ——4——
    opt.load_state_dict(opt_stat) # ——5——
    ddp_model = DDP(model, device_ids=[rank])# ——6
    inp = torch.tensor([[1.,2]]).to(rank) # ——7——
    labels = torch.tensor([[5.]]).to(rank)
    outp = ddp_model(inp)
    loss = torch.mean((outp - labels)**2)
    opt.zero_grad()
    loss.backward() # ——8——
    opt.step() # ——9
    if rank == 0:# ——10——
        torch.save(model.state_dict(), 'model_weight')
        torch.save(opt.state_dict(), 'opt_weight')
if __name__=="__main__":
    os.environ["MASTER_ADDR"] = "localhost"# ——11——
    os.environ["MASTER_PORT"] = "29500"
    world_size = 2
    mp.spawn(example, args=(world_size,), nprocs=world_size, join=True) # ——12——

以上代码包含模型在多GPU上读取权重、进行分布式训练、保存权重等过程。细节注释如下:

1、初始化进程组,由于使用GPU通信,后端应该写为NCCL。不过经过实验,即使错写为gloo,DDP内部也会自动使用NCCL作为通信模块。

2、由于后面使用DDP包裹模型进行训练,其内部会自动将所有rank的模型权重同步为rank 0的权重,因此我们只需在rank 0上读取模型权重即可。这是基于Pytorch版本1.12.1,低级版本似乎没有这个特性,需要在不同rank分别导入权重,则load需要传入map_location,如下面注释的两行代码所示。

3、这里创建model的优化器,而不是创建用ddp包裹后的ddp_model的优化器,是为了兼容单GPU训练,读取优化器权重更方便。

4、将优化器权重读取至该进程占用的GPU。如果没有map_location参数,load会将权重读取到原本保存它时的设备。

5、优化器获取权重。经过实验,即使权重不在优化器所在的GPU,权重也会迁移过去而不会报错。当然load直接读取到相应GPU会减少数据传输。

6、DDP包裹模型,为模型复制一个副本到相应GPU中。所有rank的模型副本会与rank 0保持一致。注意,DDP并不复制模型优化器的副本,因此各进程的优化器需要我们在初始化时保持一致。权重要么不读取,要么都读取。

7、这里开始模型的训练。数据需转移到相应的GPU设备。

8、在backward中,所有进程的模型计算梯度后,会进行平均(不是相加)。也就是说,DDP在backward函数添加了hook,所有进程的模型梯度的ring_reduce将在这里执行。这个可以通过给各进程模型分别输入不同的数据进行验证,backward后这些模型有相同的梯度,且验算的确是所有进程梯度的平均。此外,还可以验证backward函数会阻断(block)各进程使用梯度,只有当所有进程都完成backward之后,各进程才能读取和使用梯度。这保证了所有进程在梯度上的一致性。

9、各进程优化器使用梯度更新其模型副本权重。由于初始化时各进程模型、优化器权重一致,每次反向传播梯度也保持一致,则所有进程的模型在整个训练过程中都能保持一致。

10、由于所有进程权重保持一致,我们只需通过一个进程保存即可。

11、定义rank 0的IP和端口,使用mp.spawn,只需在主进程中定义即可,无需分别在子进程中定义。

12、创建子进程,传入:子进程调用的函数(该函数第一个参数必须是rank)、子进程函数的参数(除了rank参数外)、子进程数、是否等待所有子进程创建完毕再开始执行。

以上就是详解如何使用Pytorch进行多卡训练 的详细内容,更多关于Pytorch多卡训练的资料请关注我们其它相关文章!

(0)

相关推荐

  • pytorch DistributedDataParallel 多卡训练结果变差的解决方案

    DDP 数据shuffle 的设置 使用DDP要给dataloader传入sampler参数(torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=None, rank=None, shuffle=True, seed=0, drop_last=False)) . 默认shuffle=True,但按照pytorch DistributedSampler的实现: def __iter__(self) -> Ite

  • 详解如何使用Pytorch进行多卡训练

    目录 1.DP 2.DDP 2.1Pytorch分布式基础 2.2Pytorch分布式训练DEMO 当一块GPU不够用时,我们就需要使用多卡进行并行训练.其中多卡并行可分为数据并行和模型并行.具体区别如下图所示: 由于模型并行比较少用,这里只对数据并行进行记录.对于pytorch,有两种方式可以进行数据并行:数据并行(DataParallel, DP)和分布式数据并行(DistributedDataParallel, DDP). 在多卡训练的实现上,DP与DDP的思路是相似的: 1.每张卡都复制

  • 详解linux 驱动编写(sd卡驱动)

    随着sd卡的流行,sd卡在嵌入式设备上使用的场景也越来越多.那下面我们可以看一下,linux驱动框架上是怎么处理sd卡驱动的? 1.代码目录地址 drivers/mmc 2.基本结构 从mmc的代码结构可以看得出,主要分为两个部分,其中core为协议部分,host为各个soc的适配部分 host是我们需要真正关心的代码 3.以s3c为例,观察makefile obj-$(CONFIG_MMC_SDHCI_S3C) += sdhci-s3c.o ...... obj-$(CONFIG_MMC_S3

  • PyTorch快速搭建神经网络及其保存提取方法详解

    有时候我们训练了一个模型, 希望保存它下次直接使用,不需要下次再花时间去训练 ,本节我们来讲解一下PyTorch快速搭建神经网络及其保存提取方法详解 一.PyTorch快速搭建神经网络方法 先看实验代码: import torch import torch.nn.functional as F # 方法1,通过定义一个Net类来建立神经网络 class Net(torch.nn.Module): def __init__(self, n_feature, n_hidden, n_output):

  • 详解linux驱动编写(入门)

    在我离职之前,工作内容几乎不涉及到驱动方面的知识.我所要做的内容就是把客户对设备的请求拆分成一个一个的接口,调用驱动的设置进行配置就可以了.当然,至于驱动下面是怎么实现那就要根据具体情况而定了.比如说,有的驱动是芯片厂商直接写好的,假设芯片厂商提供了对应平台的sdk函数,那么驱动的工作就是对这些sdk函数进行封装就可以了,另外一种就是自己编写具体平台的驱动接口了.比如说,现在你需要编写串口.i2c.i2s.FLASH.网卡.LCD.触摸屏.USB驱动了.这个时候,你手里面除了一堆芯片手册,啥也没

  • 详解linux usb host驱动编写入门

    usb协议是一个复杂的协议,目前涉及到的版本就有usb1.0, usb2.0, usb3.0.大家如果打开kernel usb host目录,就会发现下面包含了ohci,uhci,ehci,xhci,whci等多种形式的控制器驱动.那么,对于我们这些不是很了解usb的开发人员,如何了解usb的代码结构呢? 1.代码分布 drivers/usb目录下面,host目录包括了host驱动代码,core目录包含了主要的api接口代码,而其他目录则主要是device驱动代码. 2.device驱动怎么看

  • 详解linux 摄像头驱动编写

    对于现代嵌入式设备,特别是手机来说,摄像头是很重要的一个设备.很多同学买手机,一看颜值,第二就看摄像头拍照如何.所以,从某个角度来说,摄像头是各个厂家主打的应用功能.那么,linux是如何支持摄像头的,我们可以来看一下? 1.代码目录地址 drivers/media 2.v4l2框架 目前linux上的camera都是按照v4l2框架来设计,它的地址位于drivers/media/v4l2-core 3.查看三星soc是如何支持camera的,可以查看drviers/media/platform

  • 详解linux电源管理驱动编写

    对于嵌入式设备来说,合适的电源管理,不仅可以延长电池的寿命,而且可以省电,延长设备运行时间,在提高用户体验方面有很大的好处.所以,各个soc厂家在这方面花了很多的功夫.下面,我们可以看看linux是如何处理电源管理驱动的. 1.代码目录 drivers/regulator 2.查看目录下的Kconfig文件 menuconfig REGULATOR bool "Voltage and Current Regulator Support" help Generic Voltage and

  • 对pytorch中的梯度更新方法详解

    背景 使用pytorch时,有一个yolov3的bug,我认为涉及到学习率的调整.收集到tencent yolov3和mxnet开源的yolov3,两个优化器中的学习率设置不一样,而且使用GPU数目和batch的更新也不太一样.据此,我简单的了解了下pytorch的权重梯度的更新策略,看看能否一窥究竟. 对代码说明 共三个实验,分布写在代码中的(一)(二)(三)三个地方.运行实验时注释掉其他两个 实验及其结果 实验(三): 不使用zero_grad()时,grad累加在一起,官网是使用accum

  • Anaconda+spyder+pycharm的pytorch配置详解(GPU)

    第一步 : 从清华大学开源软件镜像站下载Anaconda:https://mirrors.tuna.tsinghua.edu.cn/anaconda/archive/?C=M&O=D 安装过程中需要勾选如下图 装好后测试是否装好,先配置环境变量(可能anaconda安装好后自己就有了) 打开CMD,输入代码 conda list 回车出现包的信息则说明安装完成 打开Anaconda Navigator(桌面没有的话就点击左下角看最近添加)可以看到spyder已经下好了 第二步:下载CUDA(GP

随机推荐