pytorch分布式训练DDP(傻瓜版)

为什么要使用分布式训练

1. 单卡显存不够了！！！（核心原因）
2. 比较高级，比较快。

基本概念

• world_size：进程总数
• rank：每个进程的唯一编号
• nodes：节点数/主机数
• nprocs：当前节点的进程数/gpu数量
• gpu：当前GPU的序号

常用函数

以下都是忽略了导包的过程，直接调用函数，这些函数名都是分布式训练中常见的函数名

• init_process_group()
• mp.spawn(train/训练函数,nrpocs=进程数/gpus数量,args=(参数1,参数2,))
  启动一个节点/主机上面的所有进程/GPU开始训练模型

使用DataParrel

最小白的分布式训练方法，只需要增加一行代码，不过性能堪忧，只支持单机单卡，不推荐使用，效率甚至不如单卡。不需要掌握。

使用DDP

主流方法，一定要掌握，真正的分布式训练！本质是创建多个进程训练模型，每个进程占用一个GPU，不同进程做同样的事情，进程间只进行梯度传递。以resnet50进行MNIST数据集分类为例:

搭建模型

class MyResnet(nn.Module):
    def __init__(self,num_classes) -> None:
        super().__init__()
        self.conv=nn.Conv2d(in_channels=1,out_channels=3,padding=1,stride=1,kernel_size=3)
        self.model=resnet50(num_classes=num_classes)
    def forward(self,x):
        x=self.conv(x)
        x=self.model(x)
        return x

构建主函数

在主函数中我们需要利用os.environ['MASTER_ADDR]和os.environ['MASTER_PORT]指明0号节点的IP和端口，一般设置为localhost和29500。同时，使用函数mp.spawn创建多进程开始分布式训练，一个进程可以看成一个GPU，每个进程上执行的操作是完全一样的，我们只需要关注0号进程的代码，其他进程都是直接拷贝0号进程。0号进程你可以理解为你自己现在正在编辑的程序。

import os
import argparse
from torchvision.datasets import MNIST
from torchvision.transforms import Compose,ToTensor
from torchvision.models import resnet50
import torch.multiprocessing as mp 
if __name__=='__main__':
    # set environment
    os.environ['MASTER_ADDR']='localhost'
    os.environ['MASTER_PORT']='2950'
    
    # create dataset
    train_dataset=MNIST(root='/workplace/dataset/',
    train=True,transform=Compose([Resize((224,224)),ToTensor()]),download=True)
    test_dataset=MNIST(root='/workplace/dataset/',
    train=False,transform=Compose([Resize((224,224)),ToTensor()]),download=True)

    # create model
    model=MyResnet(num_classes=10)
    
    # parse argument
    parse=argparse.ArgumentParser('distribution setting')
    # 节点数/主机数
    parse.add_argument('-n','--nodes',default=1,type=int,help='the number of nodes/computer')
    # 一个节点/主机上面的GPU数
    parse.add_argument('-g','--gpus',default=1,type=int,help='the number of gpus per nodes')
    # 当前主机的编号，例如对于n机m卡训练，则nr∈[0,n-1]。对于单机多卡，nr只需为0。
    parse.add_argument('-nr','--nr',default=0,type=int,help='ranking within the nodes')
    parse.add_argument('--epochs',default=10,type=int,help='number of epochs')
    args=parse.parse_args()
    args.batch=32
    # 计算进程总数。进程总数实际上就是GPU总数
    args.world_size=args.nodes*args.gpus
    args.dataset=train_dataset
    args.model=model

    # 启动mp.spawn函数创建多进程，开始训练
    mp.spawn(
        train,
        nprocs=args.gpus,
        args=(args,)
    )

上面代码需要注意的地方有两个：

• 设置环境变量。我们必须使用os.environ['MASTER_ADDR]和os.environ['MASTER_PORT]指明0号节点/主机的IP和端口，一般直接按照示例代码设置就好。
• 使用mp.spawn创建多进程。
  mp.spawn是pytorch内置的多进程创建程序和python自带的Process差不多。使用mp.spwan后会启动多个进程进行分布式训练，一个进程就是一个GPU。mp.spawn的主要参数如下:
  • 训练函数:mp.spawn的第一个参数为训练函数，训练函数会在每个GPU上面单独执行。每个GPU都会有单独的模型，优化器，损失函数和Dataloader，互不干扰。一般而言训练函数的格式为train(gpu,args)。
  • nporcs:当前节点的进程数，其实就是当前节点的gpu数。
  • args:用于向训练函数train(gpu,args)传参，不过需要注意的是训练函数的第一个参数gpu会自动获取，所以我们只需要传第二个参数 args，一般我们建议使用argparse的形式传递第二个参数。

训练函数

使用mp.spawn创建多进程后，每个进程都会拥有一个完全独立的训练函数进行训练。在训练的过程中，每个进程的模型，优化器，损失函数和采样器Dataloader都是完全独立的，不同进程间是并行关系。在训练函数中我们需要调用init_group来初始化进程组，告诉GPU是第几个进程，使用哪个GPU。

def train(gpu,args):
    '''
    train process
    '''
    print(f'gpu:{gpu}')
    rank=args.nr*args.gpus+gpu
    # init process gropu
    dist.init_process_group(
        backend='nccl',
        world_size=args.world_size,
        rank=rank
    )
    
    # 声明训练器
    trainer=Trainer(
        gpu=gpu,
        rank=rank,
        world_size=args.world_size,
        model=args.model,
        dataset=args.dataset,
        loss_fn=torch.nn.CrossEntropyLoss()
    )
    
    # start train
    trainer(batch_size=args.batch,lr=0.001,epochs=args.epochs,pin_memory=False)

需要注意的地方只有一个，就是一定要记得调用init_process_group，这个函数用来告诉GPU你是第几个进程和总的进程数，如果没添加会运行报错。

• init_process_group相关参数:
  • backend：进程间通信方式。使用nccl就好了
  • world_size：总的进程数
  • rank：进程编号，当前进程是第几号进程

训练器

训练器是一个类，用来实现训练的过程，你也可以直接把这个训练器整合到训练函数中。

import torch
import torch.nn as nn
from torch.utils.data import Dataset,DataLoader
from torch.utils.data.distributed import DistributedSampler
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP  
import torch.distributed as dist
from datetime import datetime
from torch.nn import SyncBatchNorm as SynBN
from typing import Union
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
class Trainer(object):
def __init__(self,gpu:int,rank:int,world_size:int,model:nn.Module,dataset:Dataset,loss_fn:nn.Module):
    self.gpu=gpu
    self.rank=rank
    self.model=model
    self.dataset=dataset
    self.loss_fn=loss_fn
    self.world_size=world_size
def __call__(self,batch_size,lr,epochs,pin_memory=True):
    
    # load model
    device=torch.device(f'cuda:{self.gpu}' if torch.cuda.is_available() else 'cpu')
    model=self.model.to(device)
    if device.type!='cpu':
        model=model if self.world_size==1 else SynBN.convert_sync_batchnorm(model)
        model=DDP(model,device_ids=[self.gpu])
    
    # load dataset
    train_sampler=DistributedSampler(self.dataset,num_replicas=self.world_size,rank=self.rank)
    train_loader=DataLoader(self.dataset,batch_size=batch_size,shuffle=False,num_workers=0,pin_memory=pin_memory,sampler=train_sampler)
    
    # load criterion
    criterion=self.loss_fn.to(device)
    
    # create optimizer
    optimizer=torch.optim.Adam(params=model.parameters(),lr=lr)
    
    # train
    model.train()
    total_len=len(train_loader)
    mod=total_len//10+1
    st=datetime.now()
    for epoch in range(1,epochs+1):
        train_sampler.set_epoch(epoch)
        loss_total=0.0
        num=0.0
        for index,batch in enumerate(train_loader):
            x,y=batch
            x=x.to(device)
            y=y.to(device)
            y_pre=model(x)

            loss=criterion(y_pre,y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if self.rank==0 and ((index+1)%(mod)==0 or index+1==total_len):
                print('Epoch [{}/{}],Step [{}/{}],Loss:{:.4f}'.format(
                    epoch,
                    epochs,
                    index+1,
                    total_len,
                    loss.item()
                ))
        if self.rank==0 and (epoch%20==0 or epoch==epochs):
            torch.save(model.module.state_dict(),f'./model/model.pth')
    if self.rank==0:
        print('Training complete in:'+str(datetime.now()-st))

需要注意的地方有三个：

• 使用DDP封装模型
  我们需要在函数中使用pytorch自带的DDP类封装我们的模型，通过这个操作告诉pytorch我们的模型是一个分布式模型。一般直接用model=DDP(model,device_ids=[self.gpu])就可以了，device_ids是一个用于指明GPU的序号。
• 使用DistributedSampler进行采样
  分布式训练的本质是把数据等分成多份，分别在不同的GPU上面训练，训练完毕后再合成。所以在分布式训练时我们需要使用DistributedSampler封装数据，保证数据平均分到不同的GPU上面。例如有100份数据，用4张卡训练，则每张卡分到25份数据。如果不加DistributedSampler也能训练，只不过此时每张卡就是拥有全部数据了，也就是4张卡都有100份相同的数据。封装代码如下:

# num_replicas时进程总数，rank是当前的进程编号
train_sampler=DistributedSampler(self.dataset,num_replicas=self.world_size,rank=self.rank)

• 每次迭代都要打乱Sampler
  使用函数train_sampler.set_epoch(epoch)。否者，每个epoch采样的次序都是一样的。

启动

使用python main.py -n 1 -g 4 -nr 0 --epochs 100启动单机4卡训练resnet50

结果

参考文章

PyTorch分布式训练简明教程
Pytorch分布式训练