参考

pytorch官方文档
https://zhuanlan.zhihu.com/p/102697821
https://zhuanlan.zhihu.com/p/105755472
https://pytorch.org/tutorials/beginner/dist_overview.html#torch-nn-parallel-distributeddataparallel
https://pytorch.org/tutorials/intermediate/dist_tuto.html
https://github.com/pytorch/examples/tree/master/imagenet

单卡训练

pytorch使用GPU训练，这个比较基础了，贴一些官方的代码。

  1. 指定device
  2. 模型放入device
  3. 数据放入device ```python device = torch.device(“cuda:0”)

model.to(device)

这里执行的复制操作，会返回一个GPU的复制数据，而非重写my_tensor，所以需要给它命名

mytensor = my_tensor.to(device) 

  1. <a name="BJZhf"></a>
  2. ## 多卡训练
  3. <a name="2aimA"></a>
  4. ### torch.nn.DataParallel
  5. [torch.nn.DataParallel](https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html)是我之前常常使用的方式，它的优点是使用起来简单，缺点是GPU显存会非常不均衡。<br />之前训练的时候就有发现一个现象，0卡的显存特别特别大，其它卡的显存占用则很少。<br />这是因为torch.nn.DataParallel会并行input，loss却只会在第一块GPU上相加回传，这就造成了第一块GPU的负载会过高，且有网络通信问题。
  7. 核心代码是`model = nn.DataParallel(model)`
  9. 示例代码如下：
  10. ```python
  11. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  13. model = Model(input_size, output_size)
  14. if torch.cuda.device_count() > 1:
  15.     print("Let's use", torch.cuda.device_count(), "GPUs!")
  16.     # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
  17.     model = nn.DataParallel(model)
  19. model.to(device)
  22. input = data.to(device)
  23. output = model(input)


探究一些细节

假设batch_size = 30
使用双卡

  1. Let's use 2 GPUs!
  2.     In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
  3.     In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
  4. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  5.     In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
  6.     In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
  7. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  8.     In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
  9.     In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
  10. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  11.     In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2])
  12.     In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2])
  13. Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

使用三卡

  1. Let's use 3 GPUs!
  2.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  3.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  4.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  5. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  6.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  7.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  8.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  9. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  10.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  11.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  12.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  13. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  14.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  15.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  16.     In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
  17. Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

发现了什么？
是嘞，这种方法相当于将一个batch_size里的数据划分成多个部分分到各个卡上去。因此这种做法batch size应该大于gpu个数。概括来说，即DataParallel 会自动帮我们将数据切分load到相应GPU，将模型复制到相应GPU，进行正向传播计算计算loss，在第一张卡汇总loss并backward，权重算完后再分发到其它GPU上。

torch.nn.parallel.DistributedDataParallel

真正的分布式，一般叫DDP。官方推荐用这种方法替代multiprocessing或nn.DataParallel，更快且显存更均衡
讲DistributedDataParallel的官方文档比较短，因此先学习一下。

  1. import torch
  2. import torch.distributed as dist
  3. import torch.multiprocessing as mp
  4. import torch.nn as nn
  5. import torch.optim as optim
  6. from torch.nn.parallel import DistributedDataParallel as DDP
  9. def example(rank, world_size):
  10.     # create default process group
  11.     dist.init_process_group("gloo", rank=rank, world_size=world_size)
  12.     # create local model
  13.     model = nn.Linear(10, 10).to(rank)
  14.     # construct DDP model
  15.     ddp_model = DDP(model, device_ids=[rank])
  16.     # define loss function and optimizer
  17.     loss_fn = nn.MSELoss()
  18.     optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
  20.     # forward pass
  21.     outputs = ddp_model(torch.randn(20, 10).to(rank))
  22.     labels = torch.randn(20, 10).to(rank)
  23.     # backward pass
  24.     loss_fn(outputs, labels).backward()
  25.     # update parameters
  26.     optimizer.step()
  28. def main():
  29.     world_size = 2
  30.     mp.spawn(example,
  31.         args=(world_size,),
  32.         nprocs=world_size,
  33.         join=True)
  35. if __name__=="__main__":
  36.     main()

对代码的细节做一些解释。
一共有world_size个进程，通常每个进程控制一张卡来训练模型。每个进程需要有的操作包括（和正常训练流程不一样的标了红）：

  1. 通信信息的初始化
  2. 在第rank张卡（或者说第rank进程中）创建模型
  3. 定义损失（如果损失函数用了模型，也需要放到第rank张卡上去）
  4. 定义优化函数
  5. 前向传播获得输出，放到第rank张卡上
  6. 正常的反向传播和梯度更新

进程创建：import torch.multiprocessing as mp

在DDP基础上用Multiprocessing可以在多个gpu之间复制模型，每个GPU由一个进程控制。每个进程执行相同的任务，与其它所有进程通信。

通信：dist.init_process_group

DDP依赖ProcessGroup实现通信，因此在使用DDP前需要创建ProcessGroup实例

构造过程

初始时，DDP会构造本地模型，将state_dict()从rank0的进程开始分发到各个进程中，保证各个进程从同一状态开始。

接着，DDP会创建本地进程Reducer，用来在反向传播中梯度同步。Reducer为了保证通信效率，会将所有梯度按buckets的形式组织，比如4个梯度，分成2个buckets，以bucket为单位同步信息，确定梯度。bucket size可以用参数bucket_cap_mb控制。梯度在backet中是按model.parameters()的倒序排列的，这样可以和反向传播的顺序保持一致。下图能帮助理解，2个进程，4个参数，2个buckets，param0和param1反而在bucket1中。等所有的bucket0都计算完毕，就调用allreduce去计算平均梯度。（这里有反向传播的知识点，后面会再具体讲）

DDP除了会创建Reducer，还会创建自动去计算梯度的hook，每个参数对应一个hook。当梯度的状态被标记为ready时就会唤起这些hook。
[pytorch]多卡分布式训练 - 图1

前向传播过程

DDP获取输入并将它交给local model获取输出。

DDP的find_unused_parameters被设置为True时，会分析local model的输出，通过遍历模型输出中的autograd图将所有未使用的参数标记为Ready。反向传播过程中，Reducer需要等所有参数被标记为Ready后执行allreduce，因此有些等不到的参数应该直接将其标记为Ready（allreduce前面也提到哈，还是会在下一节反向传播的章节里讲）。 由于这种遍历会带来额外的开销，因此我们只应该在必要的时候设置为true。

反向传播过程

loss tensor会调用backward()函数, DDP控制不了tensor，因此上文中DDP才会创建自动去计算梯度的hook。当一个梯度被计算后，它对应的DDP hook就会被触发，DDP会将这个参数的梯度状态标记为Ready。当一个bucket中的所有梯度状态都是Ready之后，进程Reducer会启动异步进程（?可能是进程）allreduce去计算所有DDP进程中这个bucket中参数的平均梯度。当所有buckets的状态都是Ready后，Reducer会等待所有allreduce的计算结果。等allreduce全部计算完后，平均梯度将会写入所有参数的param.grad空间。等所有进程反向传播结束后，所有参数的grad空间应该保持一致。

优化过程

由于所有参数的grad空间保持一致，优化方法一致，所以优化结果也应该是一样的。

源码的简单解读

ProcessGroup

ProcessGroup.cpp实现DDP进程的通信，c10d的库提供了三种实现方法，分别是ProcessGroupGloo, ProcessGroupNCCL, and ProcessGroupMPI
其中的broadcast()在初始化时从rand0的进程开始分发模型的state
allreduce用来计算所有进程的梯度

接口解读

接口文档的笔记整理。

首先当然是最核心的DistributedDataParallel了

  1. CLASStorch.nn.parallel.DistributedDataParallel(module, device_ids=None, output_device=None, dim=0, broadcast_buffers=True, process_group=None, bucket_cap_mb=25, find_unused_parameters=False, check_reduction=False)
  1. batch_size应该比使用的GPU数量大
  2. 调用DistributedDataParallel前应先调用torch.distributed.init_process_group()
  3. 对于有N个GPU的显卡，应该创建N个进程，每个进程控制一张卡。
  4. 能使用的backends有gloo和nccl。nccl backend是目前最推荐也是最快的backend，单节点和多节点分布式训练中均可使用
  5. 这个模块也可以用来混合精度的分布式训练，模型可以有不同的参数精度。
  6. 如果用torch.save保存了某个进程的模型，在另一个进程中用torch.load加载，需要使用map_location。
  7. 如果用nccl或gloo，需要使用一个多workers的DataLoader，需要将multiprocessing start method改为forkserver或spawn。
  8. 将模型用DDP接口调用以后，不应该再改变模型参数

接着介绍创建进程的multiprocessing.spawn, nprocs=1, join=True, daemon=False, start_method=’spawn’))接口

  1. torch.multiprocessing.spawn(fn, args=(), nprocs=1, join=True, daemon=False, start_method='spawn')

每个主机创建nprocs个进程去执行fn函数，fn函数的输入是args。
这里稍微注意一下fn函数，它的第一个参数应该是rank（int）类型，表示第几个进程。
举例说明：

  1. mp.spawn(
  2.         main_func,
  3.         args=(n_gpus, ), # 只需要传入rank之后的其它参数
  4.         nprocs=n_gpus,
  5.         join=True
  6.     )
  8. # 第一个参数rank表示是第几个进程
  9. def main_func(rank, n_gpus):
  10.     dist.init_process_group("gloo", rank=rank, world_size=world_size)
  11.     if rank = 0:
  12.         # 可以执行一些只有第一个进程执行的操作

join=True表示执行的是阻塞连接，即操作函数会一直等候结果而不会立即返回。
daemon表示是否创建守护进程

如果join=True，这个接口会返回None
如果join是False，这个接口会返回ProcessContext

然后介绍通信的初始化接口torch.distributed.init_process_group

  1. torch.distributed.init_process_group(backend, init_method=None, timeout=datetime.timedelta(0, 1800), world_size=-1, rank=-1, store=None, group_name='')
  • backend: nccl或gloo，gloo基本只支持cpu，选nccl即可
  • init_method：指定通信的URL，如果未指定它和store，默认是”env://“
  • world_size：所有参与训练的进程数
  • rank：当前进程在所有进程中的编号

代码实战

通过前面内容的讲解，应该对torch.nn.parallel.DistributedDataParallel有了大致的了解，但还是不够。接下来会通过阅读一份优秀代码进一步加深理解。代码来自pytorch官方github库的示例，用来训练ImageNet的，本身要素比较多，我会只对保留其中和分布式训练有关的部分。

导入

没啥说的，导入需要的库

  1. import argparse
  2. import os
  3. import random
  4. import shutil
  5. import time
  6. import warnings
  8. import torch
  9. import torch.nn as nn
  10. import torch.nn.parallel
  11. import torch.backends.cudnn as cudnn
  12. import torch.distributed as dist
  13. import torch.optim
  14. import torch.multiprocessing as mp
  15. import torch.utils.data
  16. import torch.utils.data.distributed
  17. import torchvision.transforms as transforms
  18. import torchvision.datasets as datasets
  19. import torchvision.models as models

创建进程

通常代码中的world_size指训练的进程数，是nproc_per_nodenode(每个节点的进程数节点数)。这里的节点数可理解为主机数，由于我的使用场景通常是单机多卡，因此node=1。
PS：对于多机多卡的操作场景我暂时无法理解，因此会标出来
这里确认了需不需要使用分布式训练
如果需要，确认进程数量
用torch.multiprocessing.spawn创建对应的进程，使用的参数包括：需要执行的函数，每个主机的进程数量，传进函数的参数

  1. parser.add_argument('--world-size', default=-1, type=int,
  2.                     help='number of nodes for distributed training')
  3. args = parser.parse_args()
  5. if args.dist_url == "env://" and args.world_size == -1:
  6.     args.world_size = int(os.environ["WORLD_SIZE"])
  8.  args.distributed = args.world_size > 1 or args.multiprocessing_distributed
  9. if args.multiprocessing_distributed:
  10.     # Since we have ngpus_per_node processes per node, the total world_size
  11.     # needs to be adjusted accordingly
  12.     # 更新后的world_size: 所有进程的数量
  13.     args.world_size = ngpus_per_node * args.world_size
  14.     # Use torch.multiprocessing.spawn to launch distributed processes: the
  15.     # main_worker process function
  16.     # nprocs：每个主机的进程数量
  17.     mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))

进程内添加训练代码

每个进程需要执行的程序，即上文中main_worker的具体内容。
main_worker的第一个参数是固定要保留的，表示当前主机使用的gpu号（或者当前主机的进程号），程序会自动填入，一般0会对应用的0卡，1对应用的1卡，诸如此类。
rank一开始表示节点数（主机数），后来表示当前进程在所有主机所有进程中的一个编号。

多进程训练需要使用dist.init_process_group实现通信。初始化时需指定一个用来通信的地址。如果是单机多卡，可以直接指定为localhost+无人使用的端口号。示例如下：

  1. port_name = str(random.randint(10000, 99999))
  2. os.environ['MASTER_ADDR'] = 'localhost'
  3. os.environ['MASTER_PORT'] = port_name # '12355'
  4. # 之后就init_process_group即可，不需要指定init_method了

有个细节可以关注一下，作者设置的batch_size是所有GPU一起训练的时候一个批次的数量。在实际训练时，整体的batch_size被划分到各个显卡上了。workers同理。（如果使用distributedsampler，batchsize设置与单卡相同即可。）

还有一个细节，如果分布式训练，单机多卡时保存模型只保存0卡的模型。多机多卡时保留每个主机上0卡的模型

这里注意一下：每个进程都要加载一下数据集、确定criterion和优化函数。

  1. parser.add_argument('--dist-url', default='tcp://224.66.41.62:23456', type=str,
  2.                     help='url used to set up distributed training')
  3. parser.add_argument('--dist-backend', default='nccl', type=str,
  4.                     help='distributed backend')
  5. parser.add_argument('--rank', default=-1, type=int,
  6.                     help='node rank for distributed training')
  7. parser.add_argument('-b', '--batch-size', default=256, type=int,
  8.                     metavar='N',
  9.                     help='mini-batch size (default: 256), this is the total '
  10.                          'batch size of all GPUs on the current node when '
  11.                          'using Data Parallel or Distributed Data Parallel')
  12. parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
  13.                     help='number of data loading workers (default: 4)')
  15. def main_worker(gpu, ngpus_per_node, args):
  16.     # gpu表示当前主机中，本进程的编号，会和使用的卡对应，也可称之为本地序号
  17.     if args.distributed:
  18.         if args.multiprocessing_distributed:
  19.             # rank表示所有进程中，本进程的编号
  20.             args.rank = args.rank * ngpus_per_node + gpu
  22.         # 用dist.init_process_group实现通信
  23.         # 通信框架用的nccl，确认通信地址，所有需要通信的进程数量，本进信在所有进程中的编号
  24.         dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
  25.                                 world_size=args.world_size, rank=args.rank)
  28.         # For multiprocessing distributed, DistributedDataParallel constructor
  29.         # should always set the single device scope, otherwise,
  30.         # DistributedDataParallel will use all available devices.
  31.         if args.gpu is not None:
  32.             # 根据当前主机的进程号指定当前主机的显卡
  33.             torch.cuda.set_device(args.gpu)
  34.             model.cuda(args.gpu)
  36.             # 如果一个进程控制一个GPU，可以把整体batch_size划分到各个GPU中。workers同理。
  37.             # When using a single GPU per process and per
  38.             # DistributedDataParallel, we need to divide the batch size
  39.             # ourselves based on the total number of GPUs we have
  40.             args.batch_size = int(args.batch_size / ngpus_per_node)
  41.             args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)
  42.             # 把model放到对应的卡上
  43.             model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
  44.         else:
  45.             model.cuda()
  46.             # DistributedDataParallel will divide and allocate batch_size to all
  47.             # available GPUs if device_ids are not set
  48.             model = torch.nn.parallel.DistributedDataParallel(model)
  50.         # 接着是一些常规的网络训练设置
  51.         # define loss function (criterion) and optimizer
  52.         criterion = nn.CrossEntropyLoss().cuda(args.gpu)
  53.         optimizer = torch.optim.SGD(model.parameters(), args.lr,
  54.                                 momentum=args.momentum,
  55.                                 weight_decay=args.weight_decay)
  57.         # 确定数据集
  58.         train_loader = torch.utils.data.DataLoader(
  59.             train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
  60.             num_workers=args.workers, pin_memory=True, sampler=train_sampler)
  62.         # 训练
  63.         for epoch in range(args.start_epoch, args.epochs):
  64.             train(train_loader, model, criterion, optimizer, epoch, args)
  66.             # 单机多卡时只保存0卡的模型
  67.             # 多机多卡时，保留每个主机上0卡的模型
  68.             if not args.multiprocessing_distributed or (args.multiprocessing_distributed
  69.                 and args.rank % ngpus_per_node == 0):
  70.                 save_checkpoint({
  71.                     'epoch': epoch + 1,
  72.                     'arch': args.arch,
  73.                     'state_dict': model.state_dict(),
  74.                     'best_acc1': best_acc1,
  75.                     'optimizer' : optimizer.state_dict(),
  76.                 }, is_best)

train函数的具体内容如下。

  1. def train(train_loader, model, criterion, optimizer, epoch, args):
  2.     for i, (images, target) in enumerate(train_loader):
  3.         output = model(images)
  4.         loss = criterion(output, target)
  5.         optimizer.zero_grad()
  6.         loss.backward()
  7.         optimizer.step()

和nn.DataParallel的区别

DistributedDataParallel为每个GPU创建一个进程，而DataParallel使用多线程。每个GPU有专用的进程可以避免python解释器中GIL导致的性能开销