??现在并行计算的策略是不同的数据分布到不同的设备中,执行相同的任务(Data parallelism)。
??它的逻辑是,我不再拆分模型,我训练的时候模型都是一整个模型。但是我将输入的数据拆分。所谓的拆分数据就是,同一个模型在不同GPU中训练一部分数据,然后再分别计算一部分数据之后,只需要将输出的数据做一个汇总,然后再反传。其架构如下:
在PyTorch框架下,CUDA的使用变得非常简单,我们只需要显式的将数据和模型通过.cuda()方法转移到GPU上就可加速我们的训练。如下:
model = Net()
model.cuda() # 模型显示转移到CUDA上
for image,label in dataloader:
# 图像和标签显示转移到CUDA上
image = image.cuda()
label = label.cuda()
先我们来看单机多卡DP,通常使用一种叫做数据并行 (Data parallelism) 的策略,即将计算任务划分成多个子任务并在多个GPU卡上同时执行这些子任务。主要使用到了nn.DataParallel函数,它的使用非常简单,一般我们只需要加几行代码即可实现
model = Net()
model.cuda() # 模型显示转移到CUDA上
if torch.cuda.device_count() > 1: # 含有多张GPU的卡
model = nn.DataParallel(model) # 单机多卡DP训练
除此之外,我们也可以指定GPU进行并行训练,一般有两种方式
model = nn.DataParallel(model, device_ids=[0,1]) # 使用第0和第1张卡进行并行训练
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"
针对每个GPU,启动一个进程,然后这些进程在最开始的时候会保持一致(模型的初始化参数也一致,每个进程拥有自己的优化器),同时在更新模型的时候,梯度传播也是完全一致的,这样就可以保证任何一个GPU上面的模型参数就是完全一致的,所以这样就不会出现DataParallel
那样显存不均衡的问题。
进程组的相关概念
GROUP:进程组,默认情况下,只有一个组,一个 job 即为一个组,也即一个 world。(当需要进行更加精细的通信时,可以通过 new_group 接口,使用 world 的子集,创建新组,用于集体通信等。)
WORLD_SIZE:表示全局进程个数。如果是多机多卡就表示机器数量,如果是单机多卡就表示 GPU 数量。
RANK:表示进程序号,用于进程间通讯,表征进程优先级。rank = 0 的主机为 master 节点。 如果是多机多卡就表示对应第几台机器,如果是单机多卡,由于一个进程内就只有一个 GPU,所以 rank 也就表示第几块 GPU。
LOCAL_RANK:表示进程内,GPU 编号,非显式参数,由 torch.distributed.launch 内部指定。例如,多机多卡中 rank = 3,local_rank = 0 表示第 3 个进程内的第 1 块 GPU。
DDP的基本用法 (代码编写流程)
- 在使用 distributed 包的任何其他函数之前,需要使用 init_process_group 初始化进程组,同时初始化 distributed 包。
- 使用 torch.nn.parallel.DistributedDataParallel 创建 分布式模型 DDP(model, device_ids=device_ids)
- 使用 torch.utils.data.distributed.DistributedSampler 创建 DataLoader
- 使用启动工具 torch.distributed.launch 在每个主机上执行一次脚本,开始训练
添加参数 --local_rank
,这里的local_rank
参数,可以理解为torch.distributed.launch
在给一个GPU创建进程的时候,给这个进程提供的GPU号,这个是程序自动给的,不需要手动在命令行中指定这个参数。
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int) # 这个参数很重要
args = parser.parse_args()
local_rank = int(os.environ["LOCAL_RANK"]) #也可以自动获取
然后在所有和GPU相关代码的前面添加如下代码,如果不写这句代码,所有的进程都默认在你使用CUDA_VISIBLE_DEVICES
参数设定的0号GPU上面启动
torch.cuda.set_device(args.local_rank) # 调整计算的位置
接下来我们得初始化backend
,也就是俗称的后端,pytorch介绍了以下后端:
经验之谈
如果是使用cpu的分布式计算, 建议使用gloo,因为表中可以看到 gloo对cpu的支持是最好的
如果使用gpu进行分布式计算, 建议使用nccl。
GPU主机
InfiniBand连接,建议使用nccl,因为它是目前唯一支持 InfiniBand 和 GPUDirect 的后端。
Ethernet连接,建议使用nccl,因为它的分布式GPU训练性能目前是最好的,特别是对于多进程单节点或多节点分布式训练。 如果在使用 nccl时遇到任何问题,可以使用gloo 作为后备选项。 (不过注意,对于 GPU,gloo 目前的运行速度比 nccl 慢。)
CPU主机
InfiniBand连接,如果启用了IP over IB,那就使用gloo,否则使用mpi
Ethernet连接,建议使用gloo,除非有不得已的理由使用mpi。
当后端选择好了之后, 我们需要设置一下网络接口, 因为多个主机之间肯定是使用网络进行交换, 那肯定就涉及到IP之类的, 对于nccl
和gloo
一般会自己寻找网络接口,不过有时候如果网卡比较多的时候,就需要自己设置,可以利用以下代码
import os
# 以下二选一, 第一个是使用gloo后端需要设置的, 第二个是使用nccl需要设置的
os.environ['GLOO_SOCKET_IFNAME'] = 'eth0'
os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'
可以通过以下操作知道自己的网络接口,输入ifconfig, 然后找到自己IP地址的就是, 一般就是em0, eth0, esp2s0之类的
我们一般还是会选择nccl后端,设置GPU之间通信使用的后端和端口:
# ps 检查nccl是否可用
# torch.distributed.is_nccl_available ()
torch.distributed.init_process_group(backend='nccl') # 选择nccl后端,初始化进程组
之后,使用 DistributedSampler
对数据集进行划分。它能帮助我们将每个 batch 划分成几个 partition,在当前进程中只需要获取和 rank 对应的那个 partition 进行训练:
# 创建Dataloader
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, sampler=train_sampler)
然后使用torch.nn.parallel.DistributedDataParallel
包装模型:
# DDP进行训练
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])
如何启动DDP
那么如何启动DDP,这不同于DP的方式,需要使用torch.distributed.launch
启动器,对于单机多卡的情况:
CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py
# nproc_per_node: 这个参数是指你使用这台服务器上面的几张显卡
有时候虽然说,可以简单使用DP,但是DDP的效率是比DP高的,所以很多时候单机多卡的情况,我们还是会去使用DDP
DP的优点:
DP的缺点:
DDP的优点:
DDP
在每次迭代中,每个进程具有自己的 optimizer
,并独立完成所有的优化步骤,进程内与一般的训练无异。
在各进程梯度计算完成之后,各进程需要将梯度进行汇总平均,然后再由 rank=0
的进程,将其 broadcast
到所有进程。之后,各进程用该梯度来独立的更新参数。而 DP
是梯度汇总到主 GPU,反向传播更新参数,再广播参数给其他的 GPU。
DDP
中由于各进程中的模型,初始参数一致 (初始时刻进行一次 broadcast
),而每次用于更新参数的梯度也一致,因此,各进程的模型参数始终保持一致。而在DP
中,全程维护一个 optimizer
,对各 GPU 上梯度进行求和,而在主 GPU 进行参数更新,之后再将模型参数 broadcast
到其他 GPU。
相较于DP,DDP传输的数据量更少,因此速度更快,效率更高。
GIL
CPython
:是用 C 语言实现 Pyhon,是目前应用最广泛的解释器。全局锁使 Python 在多线程效能上表现不佳,全局解释器锁(Global Interpreter Lock)是 Python 用于同步线程的工具,使得任何时刻仅有一个线程在执行。GIL
,消除了来自单个 Python 进程中的多个执行线程,模型副本或 GPU 的额外解释器开销和 GIL-thrashing ,因此可以减少解释器和 GIL 使用冲突。这对于严重依赖 Python runtime 的 models 而言,比如说包含 RNN 层或大量小组件的 models 而言,这尤为重要。DDP的缺点:
Dataset定义好数据的格式和数据变换形式,DataLoader用iterative的方式不断读入批次数据。
Module
类是 torch.nn
模块里提供的一个模型构造类,是所有神经网络模块的基类,我们可以继承它来定义我们想要的模型。一般我们会重载 Module
类的 __init__
函数和 forward
函数。
PyTorch在torch.nn.init
中为我们提供了常用的初始化方法
以U-Net为例子:
组成U-Net的模型块主要有如下几个部分:
每个子块内部的两次卷积(Double Convolution)
左侧模型块之间的下采样连接,即最大池化(Max pooling)
右侧模型块之间的上采样连接(Up sampling)
输出层的处理
这里的基础部件对应上一节分析的四个模型块,根据功能我们将其命名为:DoubleConv, Down, Up, OutConv
。下面给出U-Net中模型块的PyTorch 实现:
import torch
import torch.nn as nn
import torch.nn.functional as F
DoubleConv
的实现:
class DoubleConv(nn.Module):
"""(convolution => [BN] => ReLU) * 2"""
def __init__(self, in_channels, out_channels, mid_channels=None):
super().__init__()
if not mid_channels:
mid_channels = out_channels
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(mid_channels),
nn.ReLU(inplace=True),
nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.double_conv(x)
Down
下采样的实现:
class Down(nn.Module):
"""Downscaling with maxpool then double conv"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.maxpool_conv = nn.Sequential(
nn.MaxPool2d(2),
DoubleConv(in_channels, out_channels)
)
def forward(self, x):
return self.maxpool_conv(x)
up
上采样的实现:
class Up(nn.Module):
"""Upscaling then double conv"""
def __init__(self, in_channels, out_channels, bilinear=False):
super().__init__()
# if bilinear, use the normal convolutions to reduce the number of channels
if bilinear:
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
else:
self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
self.conv = DoubleConv(in_channels, out_channels)
def forward(self, x1, x2):
x1 = self.up(x1)
# input is CHW
diffY = x2.size()[2] - x1.size()[2]
diffX = x2.size()[3] - x1.size()[3]
x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2])
# if you have padding issues, see
# https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
# https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
x = torch.cat([x2, x1], dim=1)
return self.conv(x)
OutConv
输出层的处理的实现:
class OutConv(nn.Module):
def __init__(self, in_channels, out_channels):
super(OutConv, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
def forward(self, x):
return self.conv(x)
最后利用模型块组装U-Net:
class UNet(nn.Module):
def __init__(self, n_channels, n_classes, bilinear=False):
super(UNet, self).__init__()
self.n_channels = n_channels
self.n_classes = n_classes
self.bilinear = bilinear
self.inc = DoubleConv(n_channels, 64)
self.down1 = Down(64, 128)
self.down2 = Down(128, 256)
self.down3 = Down(256, 512)
factor = 2 if bilinear else 1
self.down4 = Down(512, 1024 // factor)
self.up1 = Up(1024, 512 // factor, bilinear)
self.up2 = Up(512, 256 // factor, bilinear)
self.up3 = Up(256, 128 // factor, bilinear)
self.up4 = Up(128, 64, bilinear)
self.outc = OutConv(64, n_classes)
def forward(self, x):
x1 = self.inc(x)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x = self.up1(x5, x4)
x = self.up2(x, x3)
x = self.up3(x, x2)
x = self.up4(x, x1)
logits = self.outc(x)
return logits
如何在已有模型的基础上:
PyTorch存储模型主要采用pkl,pt,pth三种格式。对于PyTorch而言,pt, pth和pkl三种数据格式均支持模型权重和整个模型的存储,因此使用上没有差别。
一个PyTorch模型主要包含两个部分:模型结构和权重。其中模型是继承nn.Module的类,权重的数据结构是一个字典(key是层名,value是权重向量)。存储也由此分为两种形式:存储整个模型(包括结构和权重),和只存储模型权重。
from torchvision import models
model = models.resnet152(pretrained=True)
save_dir = './resnet152.pth'
# 保存整个模型
torch.save(model, save_dir)
# 保存模型权重
torch.save(model.state_dict, save_dir)
PyTorch中将模型和数据放到GPU上有两种方式——.cuda()
和.to(device)
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0' # 如果是多卡改成类似0,1,2
model = model.cuda() # 单卡
model = torch.nn.DataParallel(model).cuda() # 多卡
多卡并行的模型每层的名称前多了一个“module
”
import os
import torch
from torchvision import models
os.environ['CUDA_VISIBLE_DEVICES'] = '0' #这里替换成希望使用的GPU编号
model = models.resnet152(pretrained=True)
model.cuda()
save_dir = 'resnet152.pt' #保存路径
# 保存+读取整个模型
torch.save(model, save_dir)
loaded_model = torch.load(save_dir)
loaded_model.cuda()
# 保存+读取模型权重
torch.save(model.state_dict(), save_dir)
loaded_model = models.resnet152() #注意这里需要对模型结构有定义
loaded_model.load_state_dict(torch.load(save_dir))
loaded_model.cuda()
nn.DataParallel
函数进行分布式训练设置即可import os
import torch
from torchvision import models
os.environ['CUDA_VISIBLE_DEVICES'] = '0' #这里替换成希望使用的GPU编号
model = models.resnet152(pretrained=True)
model.cuda()
# 保存+读取整个模型
torch.save(model, save_dir)
os.environ['CUDA_VISIBLE_DEVICES'] = '1,2' #这里替换成希望使用的GPU编号
loaded_model = torch.load(save_dir)
loaded_model = nn.DataParallel(loaded_model).cuda()
# 保存+读取模型权重
torch.save(model.state_dict(), save_dir)
os.environ['CUDA_VISIBLE_DEVICES'] = '1,2' #这里替换成希望使用的GPU编号
loaded_model = models.resnet152() #注意这里需要对模型结构有定义
loaded_model.load_state_dict(torch.load(save_dir))
loaded_model = nn.DataParallel(loaded_model).cuda()
这种情况下的核心问题是:如何去掉权重字典键名中的"module",以保证模型的统一性。
对于加载整个模型,直接提取模型的module属性即可:
import os
import torch
from torchvision import models
os.environ['CUDA_VISIBLE_DEVICES'] = '1,2' #这里替换成希望使用的GPU编号
model = models.resnet152(pretrained=True)
model = nn.DataParallel(model).cuda()
# 保存+读取整个模型
torch.save(model, save_dir)
os.environ['CUDA_VISIBLE_DEVICES'] = '0' #这里替换成希望使用的GPU编号
loaded_model = torch.load(save_dir).module
对于加载模型权重,有以下几种思路: 保存模型时保存模型的module属性对应的权重
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2' #这里替换成希望使用的GPU编号
import torch
from torchvision import models
save_dir = 'resnet152.pth' #保存路径
model = models.resnet152(pretrained=True)
model = nn.DataParallel(model).cuda()
# 保存权重
torch.save(model.module.state_dict(), save_dir)
#这样保存下来的模型参数就和单卡保存的模型参数一样了,可以直接加载。也是比较推荐的一种方法。 去除字典里的module麻烦,往model里添加module简单
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2' #这里替换成希望使用的GPU编号
import torch
from torchvision import models
model = models.resnet152(pretrained=True)
model = nn.DataParallel(model).cuda()
# 保存+读取模型权重
torch.save(model.state_dict(), save_dir)
os.environ['CUDA_VISIBLE_DEVICES'] = '0' #这里替换成希望使用的GPU编号
loaded_model = models.resnet152() #注意这里需要对模型结构有定义
loaded_model.load_state_dict(torch.load(save_dir))
loaded_model = nn.DataParallel(loaded_model).cuda()
loaded_model.state_dict = loaded_dict
遍历字典去除module
from collections import OrderedDict
os.environ['CUDA_VISIBLE_DEVICES'] = '0' #这里替换成希望使用的GPU编号
loaded_dict = torch.load(save_dir)
new_state_dict = OrderedDict()
for k, v in loaded_dict.items():
name = k[7:] # module字段在最前面,从第7个字符开始就可以去掉module
new_state_dict[name] = v #新字典的key值对应的value一一对应
loaded_model = models.resnet152() #注意这里需要对模型结构有定义
loaded_model.state_dict = new_state_dict
loaded_model = loaded_model.cuda()
使用replace操作去除module
loaded_model = models.resnet152()
loaded_dict = torch.load(save_dir)
loaded_model.load_state_dict({k.replace('module.', ''): v for k, v in loaded_dict.items()})
读取整个模型再使用nn.DataParallel进行分布式训练设置
这种情况很可能会造成保存的整个模型中GPU id和读取环境下设置的GPU id不符,训练时数据所在device和模型所在device不一致而报错。
读取整个模型而不使用nn.DataParallel进行分布式训练设置
这种情况可能不会报错,测试中发现程序会自动使用设备的前n个GPU进行训练(n是保存的模型使用的GPU个数)。此时如果指定的GPU个数少于n,则会报错。在这种情况下,只有保存模型时环境的device id和读取模型时环境的device id一致,程序才会按照预期在指定的GPU上进行分布式训练。
相比之下,读取模型权重,之后再使用nn.DataParallel进行分布式训练设置则没有问题。因此多卡模式下建议使用权重的方式存储和读取模型:
import os
import torch
from torchvision import models
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2' #这里替换成希望使用的GPU编号
model = models.resnet152(pretrained=True)
model = nn.DataParallel(model).cuda()
# 保存+读取模型权重,强烈建议!!
torch.save(model.state_dict(), save_dir)
loaded_model = models.resnet152() #注意这里需要对模型结构有定义
loaded_model.load_state_dict(torch.load(save_dir)))
loaded_model = nn.DataParallel(loaded_model).cuda()
如果只有保存的整个模型,也可以采用提取权重的方式构建新的模型:
# 读取整个模型
loaded_whole_model = torch.load(save_dir)
loaded_model = models.resnet152() #注意这里需要对模型结构有定义
loaded_model.state_dict = loaded_whole_model.state_dict
loaded_model = nn.DataParallel(loaded_model).cuda()
另外,上面所有对于loaded_model修改权重字典的形式都是通过赋值来实现的,在PyTorch中还可以通过"load_state_dict"函数来实现。因此在上面的所有示例中,我们使用了两种实现方式。
loaded_model.load_state_dict(loaded_dict)
在深度学习项目里,有时候我们不仅仅需要保存模型的权重,还需要保存一些其他的参数,比如训练的epoch数、训练的loss,优化器的参数,动态调整学习策略的参数等等。这些参数可以通过字典的形式保存在一个文件里,然后在读取模型时一起读取。这里我们以下方代码为例:
torch.save({
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch,
'args': args,
}, checkpoint_path)
这些参数的读取方式也是类似的:
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
epoch = checkpoint['epoch']
args = checkpoint['args']
本文列举了pytorch初学者常见的问题,大家可以按需求进行查阅,或者对自己的pytorch的基础知识进行测试。