参考
简明教程看这里 --> pytorch分布式训练 和这篇： [PyTorch]> DDP系列第一篇：入门教程 --》 详细解答了pipeline
DDP原理篇 --> DDP系列第二篇：实现原理与源代码解析 --》 主要讲 all_reduce 和 sample 的实现
减少GPU占用看这里 --> Pytorch使用DDP加载模型时出现多进程在GPU0上占用过多显存的问题 --》解答了如何先加载到cpu解决0卡显存占用过多问题
DDP模型加载和保存看这里 – > torch DDP训练-模型保存-加载问题 --》解释和解决ddp模型名被更改后如何保存加载的问题
多机多卡更多看这里 --> Pytorch多机多卡分布式训练 --》有更细致的讲解
基本概念
在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多卡训练简明Demo
这份是一份能直接跑的简明代码，推荐收藏！
################
## 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, map_location='cpu')) # 先加载到cpu
# 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
DDP训练-模型保存-加载问题
save模型的时候，和DP模式一样，有一个需要注意的点：保存的是model.module而不是model。因为model其实是DDP model，参数是被model=DDP(model)包起来的。
模型保存：
无module形式
####方法一
state = {'epoch': epoch,
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict()}
torch.save(state, 'model_path')
####方法二
torch.save(self.model.state_dict(), 'model_path')
module形式---建议模式
####方法一
state = {'epoch': epoch,
'model': model.module.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict()}
torch.save(state, 'model_path')
####方法二
torch.save(self.model.module.state_dict(), 'model_path')
#######################################################################################
模型加载：
未用含"moduel"方式保存, 导致缺失关键“key”:Missing key(s) in state_dict
############### 方法 1: add
model = torch.nn.DataParallel(model) # 加上module
model.load_state_dict(torch.load("model_path"))
############### 方法 2: remove
model.load_state_dict({k.replace('module.', ''): v for k, v in
torch.load("model_path").items()})
############### 方法 3: remove
from collections import OrderedDict
state_dict = torch.load("model_path")
new_state_dict = OrderedDict() # create new OrderedDict that does not contain `module.`
for k, v in state_dict.items():
name = k.replace('module.', '')
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
含"moduel"方式保存
####方法一
self.model.load_state_dict(torch.load("model_path")['model'])
####方法二
self.model.load_state_dict(torch.load('model_path))
其他需要注意的地方
- 保存参数
# 1. save模型的时候，和DP模式一样，有一个需要注意的点：保存的是model.module而不是model。
# 因为model其实是DDP model，参数是被`model=DDP(model)`包起来的。
# 2. 我只需要在进程0上保存一次就行了，避免多次保存重复的东西。
if dist.get_rank() == 0:
torch.save(model.module, "saved_model.ckpt")
- 要把模型和数据放在进程对应的那张卡上
- 要使用Sampler来分发训练数据，并且shuffle不设置在Dataloder中而是Sampler中，每个epoch还需要调用Sampler的set_epoch()方法。
- 训练和验证区分较大，验证一般在主进程中进行一次验证即可，不需要sampler，操作和单卡一样，之后将数据同步给其他进程。
- 在多卡时要调用模型的其他方法或者使用单卡的模式，需要用model.module来获得原始模型，同样保存参数时也保存的是model.module的参数而不是DDP包裹的
GPU0占用更多问题的解决办法
Pytorch使用DDP加载模型时出现多进程在GPU0上占用过多显存的问题，from: (https://blog.51cto.com/u_15786578/5667478)
如果map_location参数是空的，则torch.load方法会先把模型加载到CPU，然后把模型参数复制到保存它的地方（根据上文，保存模型的位置恰好是GPU 0）。
跑在GPU1上的进程在执行到torch.load方法后，会先加载模型到CPU，之后该进程顺理成章地调用GPU0，把一部分数据复制到GPU0，也就出现了前面图中的问题。
与其说是bug，倒不如说没仔细阅读文档，两种解决方法方法。
1. 将map_location指定为CPU：
def load_checkpoint(path):
#加载到CPU
checkpoint = torch.load(path,map_location='cpu')
model = Net()
model.load_state_dict(checkpoint['model'])
model = DDP(model, device_ids=[gpu])
return model
2. 将map_location指定为local_rank对应的GPU：
def load_checkpoint(path):
#加载到CPU
checkpoint = torch.load(path,map_location='cuda:{}'.format(local_rank))
model = Net()
model.load_state_dict(checkpoint['model'])
model = DDP(model, device_ids=[gpu])
return model
多机多卡模式
复习一下，master进程就是rank=0的进程。
在使用多机模式前，需要介绍两个参数：
- 通讯的address
- master_address，也就是master进程的网络地址，默认是：127.0.0.1，只能用于单机。
- 通讯的port
- master_port，也就是master进程的一个端口，要先确认这个端口没有被其他程序占用了哦。一般情况下用默认的就行，默认是：29500
## Bash运行
# 假设我们在2台机器上运行，每台可用卡数是8
# 机器1：
python -m torch.distributed.launch --nnodes=2 --node_rank=0 --nproc_per_node 8 \
--master_adderss $my_address --master_port $my_port main.py
# 机器2：
python -m torch.distributed.launch --nnodes=2 --node_rank=1 --nproc_per_node 8 \
--master_adderss $my_address --master_port $my_port main.py
小技巧
# 假设我们只用4,5,6,7号卡
CUDA_VISIBLE_DEVICES="4,5,6,7" python -m torch.distributed.launch --nproc_per_node 4 main.py
# 假如我们还有另外一个实验要跑，也就是同时跑两个不同实验。
# 这时，为避免master_port冲突，我们需要指定一个新的。这里我随便敲了一个。
CUDA_VISIBLE_DEVICES="4,5,6,7" python -m torch.distributed.launch --nproc_per_node 4 \
--master_port 53453 main.py
为什么DDP比DP要快
来自一个博客评论：
疑问：DDP在每个GPU上参数始终一致，且每次用于更新参数的梯度也一致，那岂不是每个GPU做了重复的工作？
解答：我猜测，虽然反向传播时各GPU做了重复工作，但前向没有重复工作。因为每个GPU分到的训练数据是不一样的，
从而前向计算的梯度不一样，所以最后会有一个梯度汇总平均的操作。但反向传播的梯度一样，只是反向传播
在每个GPU重复计算了，前向计算并没有重复。反向传播重复计算也是为了防止DP中在各GPU之间大量传输模型
参数造成的通信效率问题。并且在DP中只在一个GPU进行反向，其他GPU都空着的，那么虽然DDP在各个GPU重复反向（这也是为什么DDP的GPU利用率高），其实相对于DP并没有增加额外的时间，最主要是减少了各GPU间的通信问题。
总结：DDP相对于DP最根本的速度提升点在于：DDP不用每次将模型参数 broadcast 到其他GPU，以此减少通信提升效率
参数平均
参数平均则直接计算所有模型参数的平均值。参数平均的操作在优化器执行梯度下降后，这意味着其可以作为一个辅助步骤，非常的灵活，但是参数存在一定问题：
- 数学上不等价
- 各批次数据的梯度下降方向不同，不利于收敛
- 计算和通信低效，两个阶段无法重叠，所以Pytorch的DDP设计中放弃了参数平均的方式。
梯度平均
梯度平均是将各个设备的梯度求平均，然后将这个平均梯度在各个节点的模型上作更新，这样的方式一方面在数学上和本地训练完全等价，而且可以实现异步，比参数平均更加高效。
扫一扫
97
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
