pytorch自带网络_Pytorch + Apex 训练流程记录

最近尝试用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在一个进程上效率较好。

最后提一些遇到的比较常见的问题：

1. apex编译出错：
   1. 有可能是Pytorch的cuda版本和nvcc（NVIDIA (R) Cuda compiler driver）版本不一致。查看nvcc：命令行输入nvcc --version；查看Pytorch的cuda：Python下输入torch.version.cuda 查看。不对应的话建议更换，一般需要cuda 9.2 及以上，Pytorch最好升到1.1以上。
   2. 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等时就是用这种解决办法，一试一个准。
   3. 其他问题应该google就可以搜得到。
2. 数据加速：
   1. 使用训练小网络时，梯度更新都很快，数据IO的时间消耗是训练时长的瓶颈。特别是训练Mobilenet，Shufflenet，IGCV等小网络时，这时候就要通过一些手段来加速。
   2. 例子一：ImageNet预训练，可以使用Nvidia的DALI库。通过将预处理手段放在GPU上进行来加速数据IO，会占用一部分显存。值得注意的是，DALI支持Pytorch、Mxnet和tensorflow。但是Mxnet的话只能读取record文件，所以还是先将数据生成record再调用DALI，而Pytorch则可以直接调用。详细的API在文末[1]。
   3. 例子二：我们可以用data_prefetcher类来加速，通过预读文件到GPU上来加速，具体做法在[2]。值得注意的是，链接里的类定义了DataLoader的数据输出格式，并对输出的数据做归一化，注意是否需要去除。

参考：

1. DALI文档
2. data_prefetcher类

英伟达（NVIDIA）训练深度学习模型神器APEX使用指南​blog.csdn.net 英伟达APEX，多GPU分布式训练，同步Batchnorm，自动混合精度训练法宝指南​blog.csdn.net 陈瀚可：Pytorch 安装 APEX 疑难杂症解决方案​zhuanlan.zhihu.com

lbin：pytorch 1.0 分布式​zhuanlan.zhihu.com

尼古拉·瓦砾：【PyTorch】唯快不破：基于Apex的混合精度加速​zhuanlan.zhihu.com

lbin：pytorch + apex 生活变得更美好​zhuanlan.zhihu.com

尼古拉·瓦砾：【分布式训练】单机多卡的正确打开方式（三）：PyTorch​zhuanlan.zhihu.com

MrTian：给训练踩踩油门 —— Pytorch 加速数据读取​zhuanlan.zhihu.com