Pytorch的模型加速方法：Dataparallel (DP) 和 DataparallelDistributedparallel (DDP)

Dataparallel 和 DataparallelDistributed 的区别

一、Dataparallel（DP）

1.1 Dartaparallel 的使用方式

Dataparallel 的使用方式比较简单，只需要一句话即可： net = nn.Dataparallel(net, device_ids, output_device)

其中，net 就是自己定义的网络实例，device_ids就是需要使用的显卡列表，output_device 表示参数输出结果的设备，默认情况下 output_device = device_ids[0]。因此在使用时经常发现第一块卡所占用的显存会多一些。

1.2 Dataparallel 的基本原理

Dataparallel是数据分离型，其具体做法是：在前向传播过程中，输入数据会被分成多个子部分送到不同的 device 中进行计算，而网络模型则是在每个 device 上都拷贝一份，即：输入的 batch 是平均分配到每个 device 中去，而网络模型需要拷贝到每个 device 中。在反向传播过程中，每个副本积累的梯度会被累加到原始模块中，未指明 output_device 的情况下会在 device_ids[0] 上进行运算，更新好以后把权重分发到其余卡。

1.3 Dataparallel 的注意事项

运行DataParallel模块之前，并行化模块必须在device_ids [0]上具有其参数和缓冲区。在执行DataParallel之前，会首先把其模型的参数放在device_ids[0]上。举个例子，服务器是八卡的服务器，刚好前面序号是0的卡被别人占用着，于是你只能用其他的卡来，比如你用2和3号卡，如果你直接指定 device_ids=[2, 3] 的话会出现模型初始化错误，类似于module没有复制到在 device_ids[0] 上去。那么你需要在运行train之前需要添加如下两句话指定程序可见的devices，如下：

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"

当添加这两行代码后，那么 device_ids[0] 默认的就是第2号卡，你的模型也会初始化在第2号卡上了，而不会占用第0号卡了。设置上面两行代码后，那么对这个程序而言可见的只有2和3号卡，和其他的卡没有关系，这是物理上的号卡，逻辑上来说其实是对应0和1号卡，即 device_ids[0] 对应的就是第2号卡，device_ids[1] 对应的就是第3号卡。（当然你要保证上面这两行代码需要定义在下面两行代码之前：

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

1.4 Dataparallel 的优缺点

Dataparallel 的优点就是使用起来非常简单，能够使用多卡的显存来处理数据。然而其缺点是：会造成负载不均衡的情况，成为限制模型训练速度的瓶颈。

二、DataparallelDistributed（DDP）

2.1 DDP 的基本原理

DataparallelDistributed 在每次迭代中，操作系统会为每个GPU创建一个进程，每个进程具有自己的 optimizer ，并独立完成所有的优化步骤，进程内与一般的训练无异。在各进程梯度计算完成之后，各进程需要将梯度进行汇总平均，然后再由 rank=0 的进程，将其 broadcast 到所有进程。各进程用该梯度来更新参数。由于各进程中的模型，初始参数一致 (初始时刻进行一次 broadcast)，而每次用于更新参数的梯度也一致，因此，各进程的模型参数始终保持一致。而在 DataParallel 中，全程维护一个 optimizer，对各 GPU 上梯度进行求和，而在主 GPU 进行参数更新，之后再将模型参数 broadcast 到其他 GPU。相较于 DataParallel，torch.distributed 传输的数据量更少，因此速度更快，效率更高。

2.2 DDP的使用方式

DDP使用起来比DP要麻烦一些，具体想要了解其中原理的可以参考下面几篇文章：

https://blog.csdn.net/laizi_laizi/article/details/115299263

DataParallel & DistributedDataParallel分布式训练 - 知乎 (zhihu.com)

最后，参考上述文章，整理出来了下面一份可以直接跑的代码，由于个人环境不同，可能在个别环境出现不适配的情况，可以参考上述文章进行修改。

################
## main.py文件
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无关
class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
# 假设我们的数据是这个
def get_dataset():
    transform = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    my_trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
        download=True, transform=transform)
    # DDP：使用DistributedSampler，DDP帮我们把细节都封装起来了。
    #      用，就完事儿！sampler的原理，第二篇中有介绍。
    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)
    return trainloader

### 2. 初始化我们的模型、数据、各种配置  ####
# 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设备上最快、最推荐的后端

# 准备数据，要在DDP初始化之后进行
trainloader = get_dataset()

# 构造模型
model = ToyModel().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)

# DDP: 要在构造DDP model之后，才能用model初始化optimizer。
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

# 假设我们的loss是这个
loss_func = nn.CrossEntropyLoss().to(local_rank)

### 3. 网络训练  ###
model.train()
iterator = tqdm(range(100))
for epoch in iterator:
    # DDP：设置sampler的epoch，
    # DistributedSampler需要这个来指定shuffle方式，
    # 通过维持各个进程之间的相同随机数种子使不同进程能获得同样的shuffle效果。
    trainloader.sampler.set_epoch(epoch)
    # 后面这部分，则与原来完全一致了。
    for data, label in trainloader:
        data, label = data.to(local_rank), label.to(local_rank)
        optimizer.zero_grad()
        prediction = model(data)
        loss = loss_func(prediction, label)
        loss.backward()
        iterator.desc = "loss = %0.3f" % loss
        optimizer.step()
    # 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)

################
## Bash运行
# DDP: 使用torch.distributed.launch启动DDP模式
# 使用CUDA_VISIBLE_DEVICES，来决定使用哪些GPU
# CUDA_VISIBLE_DEVICES="0,1" python -m torch.distributed.launch --nproc_per_node 2 main.py

三、总结

总之Dataparellel和Distribution都是模型训练加速的一种方法。Dataparallel (支持单机多卡)，但是速度慢（主要原因是它采用parameter server 模式，一张主卡作为reducer，负载不均衡，主卡成为训练瓶颈），在主GPU上进行梯度计算和更新，再将参数给其他gpu。而DDP则使用多线程进行加速，训练速度得到了明显的提升，但是代码修改起来比较麻烦，需要不断试错积累经验。