Pytorch分布式笔记
Pytorch多GPU计算笔记
本文内容是自己关于Pytorch多GPU分布式计算的笔记供自己学习使用,内容是从相关的博客中记录而来,侵删,如果有理解不对的地方也请大家指正。
Pytorch中分布式多GPU计算的方法有三种,一种是DP(from torch.nn import DataParallel),一种是DDP(
from torch.nn.parallel import DistributedDataParallel)。
还有一种是利用外部库APEX。
DP和DDP的区别
DP和DDP两者都能将模型和数据加载到多块GPU上,实现数据并行训练。
两者最大的区别在于进程的数目和GPU间通信的内容,DDP可以说是对DP的改进版。
需要注意的是DP中的batch_size是所有GPU中的数目,在DDP中是一个显卡的数目。
DP
详细来说,DP仅能进行单机多卡训练,而且其为单进程,仅维护了一个Optimizer。在一个batch的训练中,batch_size/GPU_nums 规模的训练数据和模型分发到每个GPU上进行损失和梯度的计算,然后将梯度值传递到GPU0上进行加和求平均得到最终梯度,Optimizer根据最终梯度对GPU0上模型的参数进行更新,并将更新后的参数分发到其他GPU中。
从DP的运行原理可以看出,其单进程控制的,所以存在着PIL(全局解释器锁)的问题,而且在每个batch中,GPU之间都要互相传模型参数,GPU的通信量大,利用效率低且负载不均衡,拥有Optimizer的GPU0一般会负载较大。
DDP
DDP可以实现单机多卡也可以实现多机多卡,为多进程,一般一个进程对应一个GPU。
与DP相比,其每个进程内都维护了一个optimizer,在一个batch的训练中,梯度计算过程与DP一样,但是DDP仅需将最终梯度分发到各GPU中,各GPU中模型利用自身的Optimizer进行参数更新,大大降低了传输的数据量并消除了负载不均衡问题。
这里需要注意的是,各GPU上的模型参数初始化必须一致,这样才能保证在进行相同的梯度下降后参数值一致,因此需要在每一个进程设置相同的随机种子,以便所有模型权重都初始化为相同的值。
Apex
APEX是英伟达开源的,完美支持PyTorch框架,用于改变数据格式来减小模型显存占用的工具。其中最有价值的是amp(Automatic Mixed Precision),将模型的大部分操作都用Float16数据类型测试,一些特别操作仍然使用Float32。并且用户仅仅通过三行代码即可完美将自己的训练代码迁移到该模型。实验证明,使用Float16作为大部分操作的数据类型,并没有降低太多效果,在一些实验中,反而由于可以增大Batch size,带来精度上的提升,以及训练速度上的提升。
Apex核心的为apex.amp和 apex.parallel.DistributedDataParallel,其中amp作用是实现混合精度训练,DistributedDataParallel就是torch中DistributedDataParallel的另一种版本,作用也是实现并行化训练。
amp的使用
在混合精度训练上,Apex 的封装十分优雅。直接使用 amp.initialize 包装模型和优化器,apex 就会自动帮助我们管理模型参数和优化器的精度了,根据精度需求不同可以传入其他配置参数。
from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
其中 opt_level 为精度的优化设置,O0(第一个字母是大写字母O):
O0:纯FP32训练,可以作为accuracy的baseline;
O1:混合精度训练(推荐使用),根据黑白名单自动决定使用FP16(GEMM, 卷积)还是FP32(Softmax)进行计算。
O2:“几乎FP16”混合精度训练,不存在黑白名单,除了Batch norm,几乎都是用FP16计算。
O3:纯FP16训练,很不稳定,但是可以作为speed的baselin
溢出问题:
因为Float16保存数据位数少了,能保存数据的上限和下限的绝对值也小了。如果我们在处理分割类问题,需要用到一些涉及到求和的操作,如sigmoid,softmax,这些操作都涉及到求和。分割问题特征图都很大,求个sigmoid可能会导致数据溢出,得到错误的结果。所以针对这些操作,仍然使用float32作为数据格式。
因为针对溢出问题,我们需要对模型进行适当的设施,如下:
from apex import amp
class xxxNet(Module):
def __init__(using_map=False)
...
...
if using_amp:
amp.register_float_function(torch, 'sigmoid')
amp.register_float_function(torch, 'softmax')
apex.parallel.DistributedDataParallel的使用
该模块的使用和torch中的DDP的使用几乎一致,只需要进行一定的修改即可:
from apex import amp
from apex.parallel import DistributedDataParallel
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
model = DistributedDataParallel(model, delay_allreduce=True)
# 反向传播时需要调用 amp.scale_loss,用于根据loss值自动对精度进行缩放
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
即需要对model和optimizer使用apm改变精度,并使用amp.scale_loss根据loss值自动对精度进行缩放,除此之外在使用时就将apex.parallel.DistributedDataParallel看作torch的DDP即可,无任何不同。
DP的使用
DP的使用十分简单,相较于单卡训练,只需要改变一点代码,将模型输入函数中即可。
model = MODEL() # 创建模型
model = model.cuda()
num_gpus = torch.cuda.device_count()
if num_gpus > 1: # 保证有多张显卡,否则没有意义
logger.info('use {} gpus!'.format(num_gpus))
model = nn.DataParallel(model)
DDP的使用
DDP的使用方法要复杂的多,这里主要记录一个进程对应一张显卡的使用方式。
相关概念
相关参数
nnode:节点数,即主机数
node_rank:节点的序号,从0开始,用来区分不同的节点
nproc_per_node:每个节点上的GPU数目
Rank:os.environ[“RANK”],对于一组能够互相发消息的进程,我们需要区分每一个进程,因此每个进程会被分配一个序号,称作rank。进程间可以通过指定rank来进行通信。
LOCAL_RANK:os.environ[“LOCAL_RANK”],每个进程在自己主机中的序号,从0开始
WORLD_SIZE:os.environ[“WORLD_SIZE”],World可以认为是一个集合,由一组能够互相发消息的进程组成,其size代表进程的总数。
master_addr:master节点的ip地址,也就是0号主机的IP地址,该参数是为了让 其他节点 知道0号节点的位,来将自己训练的参数传送过去处理
master_port:master节点的port号,在不同的节点上master_addr和master_port的设置是一样的,用来进行通信
相关函数
mp.spawn(fn, nprocs,args)
该函数的作用是生成使用 args参数运行fn的nprocs个进程。
相关参数:
fn:第一个参数是一个函数,每个进程要运行的函数。
nprocs: 第二个参数是开启的进程数目。
args:第三个参数是fn的函数实参,需要注意的是fn函数的第一个参数是进程的id,不用写入args中,spawn好像会自动分配。
torch.distributed.init_process_group(backend, init_method, timeout, world_size, rank, store, group_name)
该函数的作用是初始化默认的分布式进程组,这也将初始化分布式库。
相关主要参数:
backend:进程间通信的后端,一般选nccl。
init_method:指定如何初始化进程组的URL,默认为“env://”,即从环境变量中读取。
world_size:参与作业的进程数。
rank:当前进程的序号。
spawn函数启动
这里介绍使用spawn函数启动的方式,一个进程对应一张GPU。
1、导入分布式训练相关的模块以及定义一些相关的参数
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
2、定义并行训练函数
def example(local_rank, node_rank,nprocs_per_node, world_size):
# 初始化
torch.manual_seed(0)
torch.cuda.set_device(local_rank)
cudnn.benchmark = True
rank = local_rank + node_rank*nprocs_per_node #rank需要根据node以及GPU的数量计算
dist.init_process_group("nccl", rank=rank, world_size=world_size) #在单击多卡中,local_rank 就等于rank
# 创建模型
model = EfficientNet.from_pretrained('efficientnet-b5', num_classes=args.nclass)
.cuda()#将模型放入对应GPU中
# 将模型封装放入DDP
ddp_model = DDP(model, device_ids=[local_rank],output_device=local_rank) #要表明当前GPU
loss_fn = nn.MSELoss()
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
# 加载数据,并使用DistributedSample采样就一batch的数据平分后分发到每个GPU上,注意此时shuffle必须为false
train_dataset = ImageFolder(traindir, transform=transform_train)
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, shuffle=True)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=False, drop_last=True, pin_memory=True, num_workers=16, sampler=train_sampler)
for i in range(100):
for batch in train_dataloader:
input, label = batch[:2]
input = input.cuda()
label = label.cuda()
optimizer.zero_grad()
output = model(input)
loss = loss_fn(label, output)
loss.backward()
optimizer.step()
if rank == 0: #模型只需要一个进程保存就可
save_checkpoint(args, state_dict={
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
})
3、创建主函数
def main():
world_size = 2
parser = argparse.ArgumentParser()
parser.add_argument("--world_size", type=int)
parser.add_argument("--node_rank", type=int)
parser.add_argument("--master_addr", default="127.0.0.1", type=str)
parser.add_argument("--master_port", default="12355", type=str)
args = parser.parse_args()
local_size = torch.cuda.device_count()
#注意spawn函数可自动分配local_rank,args不需要指定,在当前主机上的spawn函数只需运行local_size个进程
mp.spawn(example,args=(args.node_rank, local_size, args.world_size,),
nprocs=local_size,join=True)
if __name__=="__main__":
main()
4、终端运行
设置为一个机器4张GPU,主节点的IP地址为192.168.0.1, 占用的端口22335,
单机多卡:
python main.py --world_size=4 --node_rank=0 --master_addr="192.168.0.1" --master_port=22335
多机多卡:
# 节点1
python main.py --world_size=8 --node_rank=0 --master_addr="192.168.0.1" --master_port=22335
# 节点2
python main.py --world_size=8 --node_rank=1 --master_addr="192.168.0.1" --master_port=22335
总结来说,相对普通训练而言,使用swape启动的DDP分布式训练主要改变在于以下几点:
1.将训练过程封装在函数中,在此函数中首需要 使用dist.init_process_group初始化进程。
2.在函数中将模型分发到相应的GPU中并使用DDP进行封装。
3.数据集的加载必须使用分布式采样器 torch.utils.data.distributed.DistributedSampler。
4.设置主函数,在主函数中使用mp.spawn开启多线程,一个线程对应一个GPU,并在每个线程中运行训练函数。
launch启动
torch.distributed.launch实际上主要完成的工作:
1、参数定义与传递。解析环境变量,并将变量传递到子进程中。
2、起多进程,调用subprocess.Popen启动多进程。
注意的地方有:
1、需要添加一个解析 local_rank的参数:
parser.add_argument(“–local_rank”, type=int)
2、dist初始化的方式 int_method取env: dist.init_process_group(“gloo”,nit_method=‘env://’)
3、DDP的设备都需要指定local_rank :
net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank], output_device=args.local_rank)
1、导入分布式训练相关的模块以及定义一些相关的参数
import torch
import torchvision
import torch.utils.data.distributed
import argparse
import torch.distributed as dist
from torchvision import transforms
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int) # 增加local_rank
args = parser.parse_args()
torch.cuda.set_device(args.local_rank)
2、定义主函数
import torch
import torchvision
import torch.utils.data.distributed
import argparse
import torch.distributed as dist
from torchvision import transforms
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int) # 增加local_rank,貌似也是自动分配
args = parser.parse_args()
torch.cuda.set_device(args.local_rank)
def main():
torch.manual_seed(0)
dist.init_process_group("nccl", init_method='env://') # init_method方式修改
trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
data_set = torchvision.datasets.MNIST('~/DATA/', train=True,
transform=trans, target_transform=None, download=True)
data_loader_train = torch.utils.data.DataLoader(dataset=data_set,
batch_size=256,
sampler=train_sampler,
num_workers=16,
pin_memory=True)
net = torchvision.models.resnet101(num_classes=10)
net.conv1 = torch.nn.Conv1d(1, 64, (7, 7), (2, 2), (3, 3), bias=False)
net = net.cuda()
# DDP 输出方式修改:
net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank],
output_device=args.local_rank)
criterion = torch.nn.CrossEntropyLoss()
opt = torch.optim.Adam(net.parameters(), lr=0.001)
for epoch in range(1):
for i, data in enumerate(data_loader_train):
images, labels = data
# 要将数据送入指定的对应的gpu中
images.to(args.local_rank, non_blocking=True)
labels.to(args.local_rank, non_blocking=True)
opt.zero_grad()
outputs = net(images)
loss = criterion(outputs, labels)
loss.backward()
opt.step()
if i % 10 == 0:
print("loss: {}".format(loss.item()))
if __name__ == "__main__":
main()
3、终端运行
假设一共有两台机器(节点1和节点2),每个节点上有8张卡,节点1的IP地址为192.168.1.1 占用的端口12355
>>> #节点1
>>>python -m torch.distributed.launch --nproc_per_node=8
--nnodes=2 --node_rank=0 --master_addr="192.168.1.1"
--master_port=12355 MNIST.py
>>> #节点2
>>>python -m torch.distributed.launch --nproc_per_node=8
--nnodes=2 --node_rank=1 --master_addr="192.168.1.1"
--master_port=12355 MNIST.py
单机多卡只需要将nnodes变为1 即可。
总的来说,langch启动的方式是对swape类似启动的封装,其相关参数需要必须在langch命令后给出,直接设定main函数即可,main函数就相当于是example函数。
单个进程占用多张卡
代码基本上一致,需要修改的地方为:
- dist.init_process_group里面的rank等于节点编号;
- world_size等于节点的总数量,不是GPU的数目;
- DDP不需要指定device。
需要注意的时,这种情况下,batchsize是单卡的数目乘每张卡上的图像数,和DP类似,和DDP不一样
##apex