Pytorch DDP原理及配置[最大限度提高GPU利用率]

前言：

pytorch在单机多卡，多机多卡之间已经做了进一步的优化，最早之前可以使用

net = torch.nn.DataParallel(net.cuda(), device_ids=range(torch.cuda.device_count()))来使用多显卡pytorch 多GPU训练总结（DataParallel的使用）_两只蜡笔的小新的博客-CSDN博客

加速运行，但是会受到python GIL锁的限制，Python GIL锁_两只蜡笔的小新的博客-CSDN博客导致对CPU的使用率不高，进而导致多显卡之间的协调不足，出现利用率较低的情况。为了解决这个微博替，需要将原来的代码进行说为的修改，使用DDP的方式提高多卡的利用率。

基本原理：

关于python线程的问题情况这篇博客Python GIL锁_两只蜡笔的小新的博客-CSDN博客，大致说的是，多线程会受到GIL锁的能影响，其实质上是一个时间段只有一个线程得到GIL，系统均分每个线程获得的时间，交替实现多线程，这是pytorch传统的DP训练模式。

而DDP，则是通过创建进程来协调多GPU之间的关系，每个进程之间的GIL锁不互相影响，pytorch在这方面已经做了api接口，较方便的可以使用【其实通过创建多进程的方式也存在很多诟病，这里不一一介绍，详细参见上篇博客】

下面介绍一些PyTorch分布式编程的基础概念。【[原创][深度][PyTorch] DDP系列第一篇：入门教程 - 知乎】

基本概念

在16张显卡，16的并行数下，DDP会同时启动16个进程。下面介绍一些分布式的概念。

group

即进程组。默认情况下，只有一个组。这个可以先不管，一直用默认的就行。

world size

表示全局的并行数，简单来讲，就是2x8=16。

# 获取world size，在不同进程里都是一样的，得到16
torch.distributed.get_world_size()

rank

表现当前进程的序号，用于进程间通讯。对于16的world sizel来说，就是0,1,2,…,15。
注意：rank=0的进程就是master进程。

# 获取rank，每个进程都有自己的序号，各不相同
torch.distributed.get_rank()

local_rank

又一个序号。这是每台机子上的进程的序号。机器一上有0,1,2,3,4,5,6,7，机器二上也有0,1,2,3,4,5,6,7

# 获取local_rank。一般情况下，你需要用这个local_rank来手动设置当前模型是跑在当前机器的哪块GPU上面的。
torch.distributed.local_rank()

修改你的代码，让它运行在DDP模式

1.修改模型加载方式

import argparse
from tqdm import tqdm
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
# 新增：
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP


### 1. 初始化我们的模型、数据、各种配置  ####
# DDP：从外部得到local_rank参数
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1, type=int)
FLAGS = parser.parse_args()
local_rank = FLAGS.local_rank

# DDP：DDP backend初始化
torch.cuda.set_device(local_rank)
dist.init_process_group(backend='nccl')  # nccl是GPU设备上最快、最推荐的后端

# 构造模型
model = yourmodel().to(local_rank)
# DDP: Load模型要在构造DDP模型之前，且只需要在master上加载就行了。
ckpt_path = None
if dist.get_rank() == 0 and ckpt_path is not None:
    model.load_state_dict(torch.load(ckpt_path))
# DDP: 构造DDP model
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

2. 修改dataloader数据加载方式

    train_sampler = torch.utils.data.distributed.DistributedSampler(my_trainset)
    # DDP：需要注意的是，这里的batch_size指的是每个进程下的batch_size。
    #      也就是说，总batch_size是这里的batch_size再乘以并行数(world_size)。
    trainloader = torch.utils.data.DataLoader(my_trainset, 
        batch_size=16, num_workers=2, sampler=train_sampler)

注意：上面数据的初始化方法的两句代码，要在DDP初始化之后进行，【dist.init_process_group(backend='nccl')】这一句之后。

3.保存模型参数

    # DDP:
    # 1. save模型的时候，和DP模式一样，有一个需要注意的点：保存的是model.module而不是model。
    #    因为model其实是DDP model，参数是被`model=DDP(model)`包起来的。
    # 2. 只需要在进程0上保存一次就行了，避免多次保存重复的东西。
    if dist.get_rank() == 0:
        torch.save(model.module.state_dict(), "%d.ckpt" % epoch)

4.终端运行启动方式

## Bash运行
# 假设我们只在一台机器上运行，可用卡数是8
python -m torch.distributed.launch --nproc_per_node 8 main.py

[原创][深度][PyTorch] DDP系列第一篇：入门教程 - 知乎

[原创][深度][PyTorch] DDP系列第二篇：实现原理与源代码解析 - 知乎
[原创][深度][PyTorch] DDP系列第三篇：实战与技巧 - 知乎