[DistributedDataParallel]分布式训练

概述

PyTorch官网教程上提供了分布式训练文档,不过说的不是很清楚。在网上找了另外一个参考:Distributed data parallel training in Pytorch,里面并没有介绍具体的实现原理,不过给出了清晰的实例。还参考了以下资源:

示例位于: zjykzj/pytorch-distributed ,完整实现可参考:ZJCV/TSN

原理

共有两种类型的分布式训练:

  • 模型分布
  • 数据分布

PyTorch提供了两种类型的实现,不过当前使用的是数据分布实现,PyTorch实现了两种方式:

  • DataParallel
  • DistributedDataParallel

前者使用单进程模式,虽然实现很简单,但是其将所有GPU的运算结果汇集到主GPU进行计算,会出现多GPU负载不均衡问题;后者对每个GPU执行单个进程,所有进程在各自GPU内计算完成loss后,将loss汇集到主GPU进行计算,大大减缓了负载不均衡的问题。当前使用的是DistributedDataParallel

简单的说,DistributedDataParallel使用了All-Reduce技术:

  • 首先每个GPU各有一个进程,每个进程里面有自己的数据加载器、模型、优化器和学习率调度器
  • 首先所有GPU各自加载同一个模型,所有模型对象进行相同的初始化
  • 将数据集按GPU个数进行切分,每个GPU载入不同的数据进行前向计算和梯度计算
  • 通过ring allreduce方式,GPU之间相互通信以平均由不同GPU计算的梯度,将平均梯度应用于各自GPU模型的反向传播以获得新权重

通过平均梯度方式,保证每个GPU上的模型参数在整个训练过程中保持一致,分布式训练相当于大批量的单GPU训练

复现

参考:

REPRODUCIBILITY

What is the differenc between cudnn.deterministic and .cudnn.benchmark?

为了保证复现性,在训练时可以设置

import numpy as np
import torch

torch.manual_seed(0)
np.random.seed(0)

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

在测试时

torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True

额外阅读