您的位置:首页 > 其它

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

2021-07-16 15:51 337 查看

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

内容来自用户分享和网络整理,不保证内容的准确性,如有侵权内容,可联系管理员处理 点击这里给我发消息
标签: