pytorch并行训练

2020-04-10 本文已影响0人 carry_xz

参考：
https://zhuanlan.zhihu.com/p/105755472

torch.nn.DataParallel 与torch.distributed 的不同

distributed较新，每个进程对应一个独立的训练过程，只有梯度等少量数据在进程中交互传递。
在各进程梯度计算完成之后，各进程需要将梯度进行汇总平均，然后再由 rank=0 的进程，将其 broadcast 到所有进程。之后，各进程用该梯度来更新参数。
由于各进程中的模型，初始参数一致 (初始时刻进行一次 broadcast)，而每次用于更新参数的梯度也一致，因此，各进程的模型参数始终保持一致。
DataParallel 中，全程维护一个 optimizer，对各 GPU 上梯度进行求和，而在主 GPU 进行参数更新，之后再将模型参数 broadcast 到其他 GPU。传输参数+梯度数据量大，等待时间长。
distributed 进程包含独立的进程锁，因此可以减少解释器和 GIL 使用冲突，这对于严重依赖 Python runtime 的 models 而言，比如说包含 RNN 层或大量小组件的 models 而言，这尤为重要。

RingAllReduce 与 TreeAllReduce

TreeAllReduce 采用 PS(参数服务器) 计算模型的分布式，gpu多时通常会遇到网络的问题。由于仅使用某一个 GPU 做服务器，该GPU需要接收其他所有 GPU 的梯度，并求平均以及 broadcast 回去，若 GPU 数量越大时，通信成本也就越高。
RingAllreduce GPU 集群被组织成一个逻辑环，每个 GPU 只从左邻居接受数据、并发送数据给右邻居，即每次同步每个 gpu 只获得部分梯度更新，等一个完整的 Ring 完成，每个 GPU 都获得了完整的参数。算法的每次通信成本是恒定的，与系统中 gpu 的数量无关，完全由系统中 gpu 之间最慢的连接决定。

分布式使用流程

并行训练又分为数据并行 (Data Parallelism) 和 模型并行两种。
数据并行指的是，多张 GPU 使用相同的模型副本，但是使用不同的数据批进行训练。而模型并行指的是，多张GPU 分别训练模型的不同部分，使用同一批数据。

基本概念

group：
即进程组。默认情况下，只有一个组，一个 job 即为一个组，也即一个 world。

当需要进行更加精细的通信时，可以通过 new_group 接口，使用 word 的子集，创建新组，用于集体通信等。

world size：
表示全局进程个数，进程总数。
rank：
表示进程序号，用于进程间通讯，表征进程优先级。rank = 0 的主机为 master 节点。
local_rank：
进程内，GPU 编号，非显式参数，由 torch.distributed.launch 内部指定。比方说， rank = 3，local_rank = 0 表示第 3 个进程内的第 1 块 GPU。

数据并行

import torch
net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])

nn.DataParallel使用起来更加简单（通常只要封装模型然后跑训练代码就ok了）。

模型的权重都是在一个进程上先算出来然后再把他们分发到每个GPU上，所以网络通信就成为了一个瓶颈，而GPU使用率也通常很低。
需要所有的GPU都在一个节点（一台机器）上，且并不支持 [Apex] 的 [混合精度训练]
容易出现GPU负载不均衡，一个GPU占满，其他占用很少。

模型并行

nn.DistributedDataParallel 进行Multiprocessing可以在多个gpu之间复制该模型，每个gpu由一个进程控制。这些GPU可以位于同一个节点上，也可以分布在多个节点上。每个进程都执行相同的任务，并且每个进程与所有其他进程通信。

DistributedDataParallel

torch.nn.parallel.DistributedDataParallel 与 apex.parallel.DistributedDataParallel 区别在于初始化方式不一样，基本可以等价替换。

from apex.parallel import DistributedDataParallel as DDP
model = DDP(model)

from torch.nn.parallel import DistributedDataParallel as DDP
model = DDP(model,device_ids=[local_rank],output_device=local_rank)

带fp16的混合精度、多gpu训练最小实例

import os
from datetime import datetime
import argparse
import torchvision
import torchvision.transforms as transforms
import torch
import torch.nn as nn
import torch.distributed as dist
# from apex.parallel import DistributedDataParallel as DDP
from torch.nn.parallel import DistributedDataParallel as DDP
from apex import amp


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-n', '--nodes', default=1, type=int, metavar='N',
                        help='number of data loading workers (default: 4)')
    parser.add_argument('-g', '--gpus', default=1, type=int,
                        help='number of gpus per node')
    parser.add_argument('-nr', '--nr', default=0, type=int,
                        help='ranking within the nodes')
    parser.add_argument('--epochs', default=10, type=int, metavar='N',
                        help='number of total epochs to run')
    parser.add_argument('--local_rank', default=-1, type=int,
                    help='node rank for distributed training')  # 注意此参数必须加上
    args = parser.parse_args()
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '8889' # 多个程序同一机器调试时需要改变port，一个程序时，这两行可以删除
    dist.init_process_group(backend='nccl')  # 初始化
    torch.cuda.set_device(args.local_rank)
    train(args)

class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc = nn.Linear(7*7*32, num_classes)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out


def train(args):
    gpu = args.local_rank # 可以用来控制模型保存，打印等
    torch.manual_seed(0)
    model = ConvNet()
    model.cuda()
    batch_size = 100
    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss().cuda()
    optimizer = torch.optim.SGD(model.parameters(), 1e-4)
    # Wrap the model
    model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
    print('Error args.local_rank:',args.local_rank)
    # model = DDP(model) # apex 两种方式都可以，注意对应关系
    model = DDP(model,device_ids=[args.local_rank],output_device=args.local_rank) # torch
    # Data loading code
    train_dataset = torchvision.datasets.MNIST(
        root='./data',
        train=True,
        transform=transforms.ToTensor(),
        download=True
    )
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        train_dataset) # 数据分割，加速训练
    train_loader = torch.utils.data.DataLoader(
        dataset=train_dataset,
        batch_size=batch_size,
        shuffle=False, # 并行时需要关闭此项
        num_workers=0,
        pin_memory=True,
        sampler=train_sampler
    )

    start = datetime.now()
    total_step = len(train_loader)
    for epoch in range(args.epochs):
        for i, (images, labels) in enumerate(train_loader):
            images = images.cuda(non_blocking=True)
            labels = labels.cuda(non_blocking=True)
            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward and optimize
            optimizer.zero_grad()
            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
            optimizer.step()
            if (i + 1) % 100 == 0 and gpu == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(
                    epoch + 1,
                    args.epochs,
                    i + 1,
                    total_step,
                    loss.item())
                )
    if gpu == 0:
        print("Training complete in: " + str(datetime.now() - start))


if __name__ == '__main__':
    main()

训练命令
python -m torch.distributed.launch --nproc_per_node=2 distributed_train.py

并行训练的模型加载

并行训练保存的模型与单卡训练的模型有些不同，会导致加载模型出问题

def _Single2Parallel(self, origin_state):
    """
    将串行的权值参数转换为并行的权值参数
    :param origin_state : 原始串行权值参数
    :return             : 并行的权值参数
    """
    converted = OrderedDict()

    for k, v in origin_state.items():
      name = "module." + k
      converted[name] = v

    return converted


def _Parallel2Single(self, origin_state):
    """
    将并行的权值参数转换为串行的权值参数
    :param origin_state : 原始串行权值参数
    :return             : 并行的权值参数
    """

    converted = OrderedDict()

    for k, v in origin_state.items():
      name = k[7:]
      converted[name] = v

    return converted