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。相较于 DataParalleltorch.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则使用多线程进行加速,训练速度得到了明显的提升,但是代码修改起来比较麻烦,需要不断试错积累经验。

posted @ 2021-07-16 15:51  ZhiboZhao  阅读(3551)  评论(0编辑  收藏  举报