pytorch多GPU分布式训练代码编写

本文主要讲述单机单卡、单机多卡的简单使用方法；

单机单卡

单机单卡就是一台机器上只有一张卡，是最简单的训练方式

对于单机单卡，我们所需要做的就是把模型和数据都拷贝到单张GPU上，但是如果GPU显存不够就会出错，这时只能调小送到GPU上的数据量或者用CPU进行训练；

在启动单机单卡的主要注意事项如下：
1）判断卡是否存在，能否顺利将数据送到GPU上；

torch.cuda.is_available()

2)将模型进行拷贝；

model.cuda()

3）数据的拷贝；

data.cuda()

4）模型的加载和保存

//模型的加载
torch.load(path, map_location= torch.device("cuda:0")) #也可以加载到cpu上
//模型的保存
torch.save（）#模型、优化器、其他变量

简单分类代码实现单机单卡

from torch import nn
import torch
import torchvision
from torch.utils.data import DataLoader
from torchvision import models

if __name__ == '__main__':
	train_data = torchvision.datasets.CIFAR10('../data',train=True,transform = transform1,download=True)
    test_data = torchvision.datasets.CIFAR10('../data',train=False,transform = transform2,download=True)
    #判断数据长度
    train_data_size = len(train_data)
    test_data_size = len(test_data)
    #加载数据集
    traindata = DataLoader(train_data,batch_size=128,pin_memory=True)
    testdata = DataLoader(test_data,batch_size=128,pin_memory=True)
    #创建网络模型
    model= models.resnet101() 
    #模型拷贝
    model= model.cuda() 
    #损失函数
    loss_fn = nn.CrossEntropyLoss()
    loss_fn = loss_fn.cuda()
    #优化器
    optim = torch.optim.SGD(cai.parameters(),lr=0.001,momentum=0.9)
    #训练次数
    total_train_steps=0
    #测试次数
    total_test_steps=0
    epoch=100
    for i in range(epoch):
        print('------第{}轮训练开始------'.format(i+1))

        #训练步骤开始
        for data in traindata:
            imgs,targets = data
            #数据拷贝
            imgs = imgs.cuda()
            targets = targets.cuda()

            outputs = model(imgs)
            loss = loss_fn(outputs,targets)
            #优化器进行调优
            optim.zero_grad()
            loss.backward()
            optim.step()

            total_train_steps =total_train_steps+1
            if total_train_steps % 100 == 0:
                print('训练次数：{}，loss：{}'.format(total_train_steps,loss.item()))

        #测试步骤开始
        model.eval()
        total_accuracy = 0
        total_test_loss = 0
        with torch.no_grad():
            for data in testdata:
                imgs, targets = data
                #数据拷贝
                imgs = imgs.cuda()
                targets = targets.cuda()
                outputs = model(imgs)
                loss = loss_fn(outputs, targets)
                total_test_loss = total_test_loss + loss.item()
                accuracy = (outputs.argmax(1) == targets).sum()
                total_accuracy = total_accuracy +accuracy
        print('整体测试集上的LOSS：{}'.format(total_test_loss ))
        print('整体测试集上的正确率：{}'.format(total_accuracy/test_data_size))
        total_test_steps = total_test_steps +1

单机多卡

单机多卡主要分为两种：

DP

1）.单进程数据并行，torch.nn.DataParallel（俗称DP）
DP的操作原理是将一个batchsize的输入数据均分到多个GPU上分别计算， 在DP模式中，总共只有一个进程（受到GIL很强限制），master节点相当于参数服务器，其会向其他卡广播其参数；在梯度反向传播后，各卡将梯度集中到master节点，master节点对搜集来的参数进行平均后更新参数，再将参数统一发送到其他卡上。这种参数更新方式，会导致master节点的计算任务、通讯量很重，从而导致网络阻塞，降低训练速度。
DP模式相较于单机单卡只需添加一行代码即可实现，model=torch.nn.DataParallel(model.cuda(), device_ids=[0,1,2,3])；

    # 创建网络模型
    model= models.resnet101() 
    #模型拷贝
    model=torch.nn.DataParallel(model.cuda() , device_ids=[0,1,2,3]) #值得注意的是，模型和数据都需要先load进GPU中，DataParallel中的module才能对其进行处理否则会报错

由于DP效率较慢，现在一般都使用的较少，一般都是使用官方推荐的DDP；

DDP

2）.多进程数据并行，torch.nn.parallel.DistributedDataParallel（俗称DDP）
DDP是每个进程控制每个GPU，与DataParallel 的单进程控制多GPU不同；DDP模式会开启N个进程，每个进程在一张显卡上加载模型，有N张卡，就会被复制N份，缓解了GIL的限制；并且DDP中每个进程在训练阶段通过Ring-Reduce的方法与其他进程通讯交换各自梯度，因此其每个进程只和自己上下游的两个进程进行通讯，极大缓解了参数服务器的通讯阻塞现象。

本文介绍一种常用的启动DDP的方式，使用pytorch官方提供的torch.distributed.launch启动器，用于在命令行执行python文件；我们需要对其进行传入相应参数：

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed training')
parser.add_argument("--gpu_id", type=str, default='0,1,2,3,4,5', help='path log files')
args = parser.parse_args()
os.environ["CUDA_DEVICE_ORDER"] = 'PCI_BUS_ID'
os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu_id

然后初始化进程 ：

torch.distributed.init_process_group("nccl"，world_size=n_gpu，rank=args.local_rank) # 第一参数nccl为GPU通信方式, world_size为当前机器GPU个数,rank为当前进程在哪个PGU上

设置进程使用第几张卡 ：

torch.cuda.set_device(args.local_rank)

对模型进行包裹：

model=torch.nn.DistributedDataParallel(model.cuda(args.local_rank), device_ids=[args.local_rank]),#这里device_ids传入一张卡即可，因为是多进程多卡，一个进程一个卡

将数据分配到不同GPU：

train_sampler = torch.util.data.distributed.DistributedSampler(train_dataset) # train_dataset为Dataset()

将train_sampler传入到DataLoader中,不需要传入shuffle=True,因为shuffle和sampler互斥 data_dataloader = DataLoader(…, sampler=train_sampler)

完整代码示例：

from torch import nn
import torch
import torchvision
from torch.utils.data import DataLoader
from torchvision import models
import argparse
import os
import torch.distributed as dist

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed training')
parser.add_argument("--gpu_id", type=str, default='0,1', help='path log files')
opt= parser.parse_args()
os.environ["CUDA_DEVICE_ORDER"] = 'PCI_BUS_ID'
os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu_id

if __name__ == '__main__':
   #初始化
    local_rank = opt.local_rank
    dist.init_process_group(backend="nccl", world_size=len(opt.gpu_id.split(',')), rank=local_rank)
    rank        = int(os.environ["RANK"])
	torch.cuda.set_device(local_rank)
	
	train_data = torchvision.datasets.CIFAR10('../data',train=True,download=True)
    test_data = torchvision.datasets.CIFAR10('../data',train=False,download=True)

    train_sampler   = torch.utils.data.distributed.DistributedSampler(train_data, shuffle=True,)
	test_sampler   = torch.utils.data.distributed.DistributedSampler(test_data, shuffle=True,)
    #判断数据长度
    train_data_size = len(train_data)
    test_data_size = len(test_data)
    #加载数据集
    traindata = DataLoader(train_data,batch_size=128,pin_memory=True,sample = train_sampler)
    testdata = DataLoader(test_data,batch_size=128,pin_memory=True,sample = test_sampler)
    #创建网络模型
    model= models.resnet101() 
    #模型拷贝
    model= model.cuda(local_rank) 
    model  = torch.nn.parallel.DistributedDataParallel(model_train, device_ids=[local_rank], find_unused_parameters=True)
    
    #损失函数
    loss_fn = nn.CrossEntropyLoss()
    loss_fn = loss_fn.cuda()
    #优化器
    optim = torch.optim.SGD(model.parameters(),lr=0.001,momentum=0.9)
    #训练次数
    total_train_steps=0
    #测试次数
    total_test_steps=0
    epoch=100
  
    for i in range(epoch):
        print('------第{}轮训练开始------'.format(i+1))
		train_sampler.set_epoch(epoch)
        #训练步骤开始
        for data in traindata:
            imgs,targets = data
            #数据拷贝
            imgs = imgs.cuda()
            targets = targets.cuda()			
            outputs = model(imgs)
            loss = loss_fn(outputs,targets)
            #优化器进行调优
            optim.zero_grad()
            loss.backward()
            optim.step()
            total_train_steps =total_train_steps+1
            if total_train_steps % 100 == 0:
                print('训练次数：{}，loss：{}'.format(total_train_steps,loss.item()))

        #测试步骤开始
        model.eval()
        total_accuracy = 0
        total_test_loss = 0
        with torch.no_grad():
            for data in testdata:
                imgs, targets = data
                #数据拷贝
                imgs = imgs.cuda()
                targets = targets.cuda()
                outputs = model(imgs)
                loss = loss_fn(outputs, targets)
                total_test_loss = total_test_loss + loss.item()
                accuracy = (outputs.argmax(1) == targets).sum()
                total_accuracy = total_accuracy +accuracy
        print('整体测试集上的LOSS：{}'.format(total_test_loss ))
        print('整体测试集上的正确率：{}'.format(total_accuracy/test_data_size))
        total_test_steps = total_test_steps +1

注意事项
1.在每个epoch开始时候，需要调用train_sampler.set_epoch(epoch)使得数据充分打乱，要不然每个epoch返回数据是相同的；
2.执行命令的时候需加入-m torch.distributed.launch参数，nproc_per_node执行进程个数/GPU个数，launch会像xx.py传入args.local_rank,local_rank从0到n_gpus - 1个索引
python -m torch.distributed.launch --nproc_per_node=n_gpus --master_port 22225 xx.py
3.launch会像train.py传入args.local_rank,local_rank从0到n_gpus - 1个索引，train.py需要接受local_rank的参数
4.注意Batch_size为每个GPU的Batch_size
5.MASTER_PORT: master节点的端口号,我们在使用服务器的时候，它的默认节点会被别人使用，所以别人在使用launch启动时会报错，所以我们需要随机生成一个端口号启动。