上手Pytorch分布式训练DDP
自己这两天改代码的一些经历，记录一下。
DDP
对于多卡训练，Pytorch支持nn.DataParallel 和nn.parallel.DistributedDataParallel这两种方式。其中nn.DataParallel 最简单但是效率不高，nn.parallel.DistributedDataParallel(DDP)不仅支持多卡，同时还支持多机分布式训练，速度更快，更加强大。理论上来说，使用2张GPU的速度应该是1张GPU训练速度的两倍。为了加速巡练，，我将原来的单卡程序改成DDP多卡程序，看到训练速度底能够有多少提升。
单机单卡
首先是单卡的训练程序，这里的Demo模型是简单的ResNet来对CIFAR10进行分类。训练程序很简单，一眼就能看明白。
训练命令
python train_single_gpu.py --device 1 --batch_size 32
"""
train_single_gpu.py
Adapted from https://github.com/wmpscc/CNN-Series-Getting-Started-and-PyTorch-Implementation
"""
import torch
import torchvision.transforms as transforms
import argparse
from torch import nn, optim
from torch.nn import functional as F
from torchvision import datasets
from tqdm import tqdm
from ResNet import ResNet
from utils import evaluate_accuracy
def main(opt):
"""
Train and valid
"""
batch_size = opt.batch_size
device = torch.device('cuda', opt.device if torch.cuda.is_available() else 'cpu')
print("Using device:", device)
#加载CIFAR10数据
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
valset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=opt.num_workers)
val_loader = torch.utils.data.DataLoader(valset, batch_size=batch_size, shuffle=False, num_workers=opt.num_workers)
#定义模型
resnet = ResNet(classes=opt.num_classes)
resnet = resnet.to(device)
#损失函数
optimizer = optim.Adam(resnet.parameters(), lr=opt.lr)
lossFN = nn.CrossEntropyLoss()
num_epochs = opt.epoch
for epoch in range(num_epochs):
sum_loss = 0
sum_acc = 0
batch_count = 0
n = 0
for X, y in tqdm(train_loader):
X = X.to(device)
y = y.to(device)
y_pred = resnet(X)
loss = lossFN(y_pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
sum_loss += loss.cpu().item()
sum_acc += (y_pred.argmax(dim=1) == y).sum().cpu().item()
n += y.shape[0]
batch_count += 1
if epoch > 0 and epoch % 2 == 0:
test_acc = evaluate_accuracy(val_loader, resnet)
print("epoch %d: loss=%.4f \t acc=%.4f \t test acc=%.4f" % (epoch + 1, sum_loss / n, sum_acc / n, test_acc))
if __name__ == '__main__':
parser = argparse.ArgumentParser('Single GPU training script.')
parser.add_argument('--batch_size', type=int, default=32)
parser.add_argument('--num_workers', type=int, default=4)
parser.add_argument('--num_classes', type=int, default=10)
parser.add_argument('--epoch', type=int, default=20)
parser.add_argument('--lr', type=float, default=0.01)
parser.add_argument('--device', type=int, default=0, help='select GPU devices')
opt = parser.parse_args()
print("opt:", opt)
main(opt)
训练结果
在CIFAR10上训练batch-size=32情况下，平均一个epoch需要1分钟，在这种小数据上都要1分钟才能跑一轮，属实有点慢，需要加速。
单机多卡
分布式训练可以分为单机多卡和多机多卡。
改成单机多卡需要不同GPU之间进行通信，多机多卡还需要不同主机之间进行通信。如果这些都要自己实现的话，我选择放弃哈哈哈。不过好在Pytorch实现了这一系列接口，方便我们将单卡程序改成多卡程序。
对于多卡训练，我理解的实际上就是使用多个进程在多个GPU上同时进行训练，在得到多个计算结果之后在汇总平均各个GPU上的梯度，然后再根据得到的梯度对各个GPU上的模型参数进行更新，从而实现加速。
这里需要明确一些概念，对于多卡训练，需要初始化一个进程组，一般来说训练涉及所有的进程都属于这个进程组。一般来说，需要定义下面这些变量来标识具体某一个线程
变量world_size表示训练使用的全局进程数，比如有2台机器，每台机器有4张GPU，每个GPU使用一个线程训练，那么world_size=2*4=8
变量rank表示某个进程在整个进程组内的序号用来唯一表示某一进程。
变量local_rank表示某一GPU，由torch.distributed.launch内部指定，所以一般需要在parser里面定义这个参数。例如，rank=2,local_rank=1表示第2个进程内的第2块GPU。
需要使用到的函数
初始化进程组 init_process_group
torch.distributed.init_process_group(backend,
init_method=None,
timeout=datetime.timedelta(0, 1800),
world_size=-1,
rank=-1,
store=None)
函数功能
该函数需要在每个进程中进行调用，用于初始化该进程，该函数必须在 distributed 内所有相关函数之前使用。
参数
backend ：指定当前进程要使用的通信后端
可选nccl,gloo,mpi,推荐nccl
init_method ： 指定当前进程组初始化方式
可选参数，默认为 env://方式，就是读取环境变量的初始化方式，还可以选择tcp初始化方式。
rank ： 指定当前进程的优先级
也就是当前进程的编号，例如rank=0表示该进程是主进程。
world_size
任务中用到的总的进程数
timeout ： 指定每个进程的超时时间
store
所有 worker 可访问的 key / value，用于交换连接 / 地址信息。与 init_method 互斥。
DistributedDataParallel
torch.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)
函数功能
对模型module进行分布式封装，并将模型分配到指定GPU上，在更新参数时，对每个GPU上的梯度进行汇总并求平均，之后在更新各个GPU上的模型。
参数
module
要包装的模型
device_ids
int 列表或 torch.device 对象，用于指定要并行的GPU。对于数据并行，即将一个完整的模型放置于一个 GPU 上（single-device module）时，需要提供该参数，表示将模型副本拷贝到哪些 GPU 上，每个GPU上有一个完整的模型。对于模型并行的情况，即一个模型，分散于多个 GPU 上的情况（multi-device module），以及 CPU 模型，该参数比必须为 None，或者为空列表。
output_device
int 或者 torch.device，表示结果输出的设备
broadcast_buffers
bool 值，默认为 True，表示在 forward() 函数开始时，对模型的 buffer 进行同步 (broadcast)。
process_group
默认为 None，表示使用由 torch.distributed.init_process_group 创建的默认进程组
DistributedSampler
torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=None, rank=None)
函数功能
对数据集进行划分并分配给不同GPU。例如使用4个GPU，那么就将数据集划分为4份，每个GPU分配其中一份，不同GPU的数据是不一样的。
参数
dataset
创建好的dataset
num_replicas
分布式训练中，参与训练的进程数
rank
当前进程的 rank 序号
初始化方式
常用的初始化方式一般有tcp初始化方式和环境变量env初始化方式
对于单机多卡这里使用环境变量初始化方式，使用起来感觉比较方便，而且torch.distributed.launch默认就是环境变量初始化。
python -m torch.distributed.launch --nproc_per_node=number_gpus train.py --args
大致使用流程
- 在使用 distributed 包的任何其他函数之前，需要使用 init_process_group 初始化进程组，同时初始化 distributed 包。
- 将定义好的单卡模型用DDP包装起来，model=DDP(model, device_ids=device_ids)
- 创建sampler，用来对数据集进行划分，将不同子数据集分配给不同GPU，实现数据并行。
- 使用 destory_process_group() 销毁进程组
- 使用pytorch自带的启动工具torch.distributed.launch（或者别的工具也行）启动训练程序
下面是修改的多卡训练程序，这里给出我的完整程序。
训练命令
python -m torch.distributed.launch --nproc_per_node=2 train_ddp_single_node.py --batch_size 32
"""
train_ddp_single_gpu.py
"""
import torch
import torchvision.transforms as transforms
import argparse
from torch import nn, optim
from torch.nn import functional as F
from torchvision import datasets
import torch.distributed as dist #1. DDP相关包
import torch.utils.data.distributed
from tqdm import tqdm
from ResNet import ResNet
from utils import evaluate_accuracy
def main(opt):
"""
Train and valid
"""
dist.init_process_group(backend='nccl', init_method=opt.init_method) #4.初始化进程组，采用nccl后端
batch_size = opt.batch_size
device = torch.device('cuda', opt.local_rank if torch.cuda.is_available() else 'cpu')
print("Using device:{}\n".format(device))
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset, shuffle=True) #5. 分配数据，将数据集划分为N份，每个GPU一份
train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, num_workers=opt.num_workers, sampler=train_sampler) #注意要loader里面要指定sampler，这样才能将数据分发到多个GPU上
nb = len(train_loader)
#6.一般只在主进程进行验证，所以在local_rank=-1或者0的时候才实例化val_loader
if opt.local_rank in [-1, 0]:
valset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
val_loader = torch.utils.data.DataLoader(valset, batch_size=batch_size, shuffle=False, num_workers=opt.num_workers)
resnet = ResNet(classes=opt.num_classes)
resnet = resnet.to(device)
resnet = torch.nn.parallel.DistributedDataParallel(resnet, device_ids=[opt.local_rank], output_device=opt.local_rank) #7. 将模型包装成分布式
optimizer = optim.Adam(resnet.parameters(), lr=opt.lr)
cross_entropy = nn.CrossEntropyLoss()
num_epochs = opt.epoch
for epoch in range(num_epochs):
if opt.local_rank != -1:
train_loader.sampler.set_epoch(epoch) #不同的epoch设置不同的随机数种子，打乱数据
loader = enumerate(train_loader)
if opt.local_rank in [-1, 0]:
loader = tqdm(loader, total=nb) #只在主进程打印进度条
sum_loss = 0
sum_acc = 0
batch_count = 0
n = 0
for _, (X, y) in loader:
X = X.to(device)
y = y.to(device)
y_pred = resnet(X)
loss = cross_entropy(y_pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
sum_loss += loss.cpu().item()
sum_acc += (y_pred.argmax(dim=1) == y).sum().cpu().item()
n += y.shape[0]
batch_count += 1
if opt.local_rank in [-1, 0] and epoch % 2 == 0 and epoch > 0:
test_acc = evaluate_accuracy(val_loader, resnet)
print("epoch %d: loss=%.4f \t acc=%.4f \t test acc=%.4f" % (epoch + 1, sum_loss / n, sum_acc / n, test_acc))
if __name__ == '__main__':
parser = argparse.ArgumentParser('DDP training script.')
parser.add_argument('--batch_size', type=int, default=16)
parser.add_argument('--num_workers', type=int, default=4)
parser.add_argument('--num_classes', type=int, default=10)
parser.add_argument('--epoch', type=int, default=20)
parser.add_argument('--lr', type=float, default=0.01)
parser.add_argument('--local_rank', type=int, default=-1, help='local_rank of current process') #2. 指定local_rank，这个参数必须要有
parser.add_argument('--init_method', default='env://') #3.指定初始化方式，这里用的是环境变量的初始化方式
opt = parser.parse_args()
if opt.local_rank in [-1, 0]:
print("opt:", opt)
main(opt)
训练结果
双卡训练
在这里可以看到，在使用2张GPU，每个GPU的batch-size=32的情况下，平均一个epoch的训练时长大致是34s左右，速度接近翻倍，说明DDP分布式训练确实可以提速很多。
用2张卡就能快近一倍，那要是4张卡岂不是不到20s就能跑完？试一试！
4卡训练
python -m torch.distributed.launch --nproc_per_node=4 train_ddp_single_node.py --batch_size 32
这里用了25s左右，比预想的慢一些，因为这里训练程序并没有优化很好并且GPU之间需要同步梯度等信息，所以4卡模式下每张张卡的速度要比单卡模式的速度慢一些，但是四张卡加起来仍然提速很多。
多机多卡
手上没有多台机器，暂时用不到。
这里贴上多机连接，需要的自取。
巨人的肩膀
https://zhuanlan.zhihu.com/p/158375055
https://zhuanlan.zhihu.com/p/158375055
https://zhuanlan.zhihu.com/p/68717029