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

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

文章插图
由于模型并行比较少用，这里只对数据并行进行记录。对于pytorch，有两种方式可以进行数据并行：数据并行（DataParallel, DP）和分布式数据并行（DistributedDataParallel, DDP）。
在多卡训练的实现上，DP与DDP的思路是相似的：
1、每张卡都复制一个有相同参数的模型副本。
2、每次迭代，每张卡分别输入不同批次数据，分别计算梯度。
3、DP与DDP的主要不同在于接下来的多卡通信：
DP的多卡交互实现在一个进程之中，它将一张卡视为主卡，维护单独模型优化器。所有卡计算完梯度后，主卡汇聚其它卡的梯度进行平均并用优化器更新模型参数，再将模型参数更新至其它卡上。
DDP则分别为每张卡创建一个进程，每个进程相应的卡上都独立维护模型和优化器。在每次每张卡计算完梯度之后，进程之间以NCLL（NVIDIA GPU通信）为通信后端，使各卡获取其它卡的梯度。各卡对获取的梯度进行平均，然后执行后续的参数更新。由于每张卡上的模型与优化器参数在初始化时就保持一致，而每次迭代的平均梯度也保持一致，那么即使没有进行参数复制，所有卡的模型参数也是保持一致的。
Pytorch官方推荐我们使用DDP 。DP经过我的实验，两块GPU甚至比一块还慢。当然不同模型可能有不同的结果。下面分别对DP和DDP进行记录。
DPPytorch的DP实现多GPU训练十分简单，只需在单GPU的基础上加一行代码即可。以下是一个DEMO的代码。
import torchfrom torch import nnfrom torch.optim import Adamfrom torch.nn.parallel import DataParallelclass 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 outpmodel = 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中计算损失，并进行反向传播计算梯度、优化器更新参数。
DDP为了对分布式编程有基本概念，首先使用pytorch内部的方法实现一个多进程程序，再使用DDP模块实现模型的分布式训练。
Pytorch分布式基础首先使用pytorch内部的方法编写一个多进程程序作为编写分布式训练的基础。
import os, torchimport torch.multiprocessing as mpimport torch.distributed as distdef 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操作以及其它通信方式请看下图：