最近尝试用Pytorch的分布式训练API进行一系列实验,踩了不少坑,也看了很多资料。在这里简单做一下总结。内容包括有:Pytorch自带的分布式训练API,以及Nvidia的apex附带的API,最后再额外写一下Nvidia的DALI数据读取与处理加速的方法。文末是一些查看后觉得有用的链接,可以先看看里面的内容,本文只是对其做一些补充。
在一些大的视觉任务例如语义分割和目标检测中由于数据IO不是瓶颈,所以DALI其实不太需要,我是在ImageNet预训练阶段会用到DALI,提速确实明显。
首先是Pytorch自带的分布式API:
import
以上是环境初始化以及数据读取部分的设置,无论使用apex还是Pytorch自带的分布式API,上面都是必须做的,以下模型读取的话Pytorch的API和apex有所不同,主要是apex不仅可以分布式训练,还可以使用混合精度模式以及同步bn(Sync-BN)。参考代码:
from apex.parallel import DistributedDataParallel as DDP
from apex.parallel import convert_syncbn_model, SyncBatchNorm
from apex.fp16_utils import *
from apex import amp, optimizers
from apex.multi_tensor_apply import multi_tensor_applier
model = XXXNet()
model.train()
model.to(device) # device是上面设置好的device。
# Sync-BatchNorm
# 使用apex自带的同步BN层。
# 调用函数后会自动遍历model的层,将Batchnorm层替换。
model = convert_syncbn_model(model)
optimizer = optim.SGD(model.parameters(), lr=.., decay=.., ...)
# 用于记录训练过程中的loss变化。
def reduce_tensor(tensor, world_size=1):
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
tensor /= world_size
return tensor
# 通过调整下面的opt_level实现半精度训练。
# opt_level选项有:'O0', 'O1', 'O2', 'O3'.
# 其中'O0'是fp32常规训练,'O1'、'O2'是fp16训练,'O3'则可以用来推断但不适合拿来训练(不稳定)
# 注意,当选用fp16模式进行训练时,keep_batchnorm默认是None,无需设置;
# scale_loss是动态模式,可以设置也可以不设置。
model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level,
keep_batchnorm_fp32=keep_bn_32, scale_loss=scale_loss)
# Pytorch自带的分布式API
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])
# Apex的分布式API
model = DDP(model)
# 训练过程中:
while training:
output = model(input)
loss = criterion(output, target)
# 方便记录损失下降
log_loss = reduce_tensor(log_loss.clone().detach_())
optimizer.zero_grad()
# 假如设置了fp16训练模式,则要将反传过程替换为以下。
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
optimizer.step()
关于模型保存和重用:
# Save checkpoint
# 注意,在保存权重时要注意判断local_rank。
# 因为是多进程任务,假如多个进程同时写同一个文件在同一个地址会损坏文件。
# 通常是local_rank 为默认值时保存权重,由于多个进程之间的保持GPU通信,所以每个进程的权重都是一样的。
if local_rank == 0:
checkpoint = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'amp': amp.state_dict()
}
torch.save(checkpoint, 'amp_checkpoint.pt')
...
# 模型重用,把权重放在‘cpu’模式上,然后读取,再设置分布式。
net = XXXNet()
net.train()
resume = 'XXXNet.pth'
checkpoint = torch.load(resume, map_location='cpu') # CPU mode
net.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
amp.load_state_dict(checkpoint['amp'])
注意,由于我自己还没有尝试过模型重用,所以具体流程还不是特别清楚,所以这里遇到什么问题可以交流下。
最后一步就是开启训练:
# 在终端输入:
NCCL_DEBUG=INFO CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py
# NCCL_DEBUG=INFO 表示打印nccl初始化信息,可选,用于debug。
# CUDA_VISIBLE_DEVICES=0,1 用来指定训练用GPU
# python -m torch.distributed.launch 必须加,用来初始化local_rank到每一个进程
# --nproc_per_node=2 通常一台机器上用多少GPU就设置多少,每个GPU在一个进程上效率较好。
最后提一些遇到的比较常见的问题:
- apex编译出错:
- 有可能是Pytorch的cuda版本和nvcc(NVIDIA (R) Cuda compiler driver)版本不一致。查看nvcc:命令行输入nvcc --version;查看Pytorch的cuda:Python下输入torch.version.cuda 查看。不对应的话建议更换,一般需要cuda 9.2 及以上,Pytorch最好升到1.1以上。
- C++文件编译出错。准确来说是Python的C++编译文件,像我的话我的g++在anaconda路径下,和C++其他标准库路径不一致,要在Python代码(通常是setup.py文件中)的Extension部分添加extra_link_args=['-L/usr/lib/x86_64-linux-gnu/']。这个问题一般是出在各种涉及到C++编译例如目标检测的NMS或者ROI操作。之前编译目标检测包如mmdetection和Simpledet等时就是用这种解决办法,一试一个准。
- 其他问题应该google就可以搜得到。
- 数据加速:
- 使用训练小网络时,梯度更新都很快,数据IO的时间消耗是训练时长的瓶颈。特别是训练Mobilenet,Shufflenet,IGCV等小网络时,这时候就要通过一些手段来加速。
- 例子一:ImageNet预训练,可以使用Nvidia的DALI库。通过将预处理手段放在GPU上进行来加速数据IO,会占用一部分显存。值得注意的是,DALI支持Pytorch、Mxnet和tensorflow。但是Mxnet的话只能读取record文件,所以还是先将数据生成record再调用DALI,而Pytorch则可以直接调用。详细的API在文末[1]。
- 例子二:我们可以用data_prefetcher类来加速,通过预读文件到GPU上来加速,具体做法在[2]。值得注意的是,链接里的类定义了DataLoader的数据输出格式,并对输出的数据做归一化,注意是否需要去除。
参考:
- DALI文档
- data_prefetcher类
![4576dbf6c00ceba4297e285cf4b278c8.png](https://img-blog.csdnimg.cn/img_convert/4576dbf6c00ceba4297e285cf4b278c8.png)
![878475625a3344ee51bd782924a0abc2.png](https://img-blog.csdnimg.cn/img_convert/878475625a3344ee51bd782924a0abc2.png)
![4f7dc7b283e7f452d60a54f0e998046b.png](https://img-blog.csdnimg.cn/img_convert/4f7dc7b283e7f452d60a54f0e998046b.png)