上手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

4卡训练

这里用了25s左右，比预想的慢一些，因为这里训练程序并没有优化很好并且GPU之间需要同步梯度等信息，所以4卡模式下每张张卡的速度要比单卡模式的速度慢一些，但是四张卡加起来仍然提速很多。

多机多卡

手上没有多台机器，暂时用不到。
这里贴上多机连接，需要的自取。

巨人的肩膀

https://zhuanlan.zhihu.com/p/158375055
https://zhuanlan.zhihu.com/p/158375055
https://zhuanlan.zhihu.com/p/68717029