使用Pytorch进行多卡训练( 二 ) _生活百科

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
Pytorch分布式训练DEMO我们实际上可以根据上面的分布式基础写一个分布式训练，但由于不知道pytorch如何实现GPU间模型梯度的求和，即官方教程中所谓的ring_reduce（没找到相关API），时间原因，就不再去搜索相关方法了。这里仅记录pytorh内部的分布式模型训练，即利用DDP模块实现。Pytorch版本1.12.1 。
import torch,osimport torch.distributed as distimport torch.multiprocessing as mpimport torch.optim as optimfrom torch.nn.parallel import DistributedDataParallel as DDPfrom torch import nndef 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])# ——6inp = 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() # ——9if 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 = 2mp.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之后，各进程才能读取和使用梯度。这保证了所有进程在梯度上的一致性。

使用Pytorch进行多卡训练( 二 )

推荐阅读

教你如何使用迅雷离线下载[新发布]

在“再见阿郎”中的刘金虎怎么又是金鹏？

如何通过女人面相看婚姻哪些面相的女人婚姻好

梦见和当官的同学在一起商人梦见当官的同学

皮蛋是生的还是熟的怎么判断皮蛋是生的还是熟的

梦见水里有黑鱼在游

一斤几克重，一斤等于多少千克

我们好像在哪见过歌词我们好像在哪见过歌词歌谱

商业承兑汇票是什么意思商业承兑汇票的理解

煤炭公司起名大全免费煤炭行业起名大全

天麻能泡白酒吗天麻可以和虫草花一起煲汤吗

etc不抬杆可以倒回吗 etc不抬杆可以倒回会被扣分吗

节日由来英语节日由来

瑜伽入门动作图解瑜伽入门动作图解上

五岳奇险天下第一山是哪座山天下第一山是哪座山

孕妇梦见和别人吵架孕妇梦见别人挣了很多钱

喝番茄汁有什么好处番茄汁的好处有哪些

忙碌了多少天辛苦了多少年是啥歌快去听一下吧

关于洛阳昊远机械设备有限公司简述洛阳昊远机械设备有限公司

“八岐之术”是什么？