• 如果您觉得本站非常有看点,那么赶紧使用Ctrl+D 收藏吧

轻量化模型训练加速的思考(Pytorch实现)

开发技术 开发技术 4周前 (09-01) 21次浏览

0. 引子

在训练轻量化模型时,经常发生的情况就是,明明 GPU 很闲,可速度就是上不去,用了多张卡并行也没有太大改善。

如果什么优化都不做,仅仅是使用nn.DataParallel这个模块,那么实测大概只能实现一点几倍的加速(按每秒处理的总图片数计算),不管用多少张卡。因为卡越多,数据传输的开销就越大,副作用就越大。

为了提高GPU服务器的资源利用率,尝试了一些加速的手段。

基于Pytorch1.6.0版本实现,官方支持amp功能,不再需要外部apex库;
此外比较重要的库是Dali。

梳理了训练框架,并将参考代码放到Github上。
如果觉得对你有所启发,请给个star呀。

参考代码

1. 训练速度的瓶颈及应对思路

这边主要说的是CV领域,但在其他领域,思路应该也是相通的。
模型训练过程中,影响整体速度的因素主要有以下几点:

  1. 将数据从磁盘读取到内存的效率;
  2. 对图片进行解码的效率;
  3. 对样本进行在线增强的效率;
  4. 网络前向/反向传播和Loss计算的效率;

针对这几个因素,分别采取如下几种应对思路:

  1. 加快数据读取可以有几种思路:
    • 采取类似TF的tfrecord或者Caffe的lmdb格式,提前将数据打包,比反复加载海量的小文件要快很多,但pytorch没有通用的数据打包方式;
    • 在初始化时,提前将所有数据加载到内存中(前提是数据集不能太大,内存能装得下);
    • 将数据放在SSD而非HDD,可以大大提速(前提是你有足够大的SSD);
  2. 提升图片解码速度,可以考虑采用NVIDIA-DALI库,能够利用GPU来加速JPG格式的图片解码,针对其他格式的图片(如PNG),不能实现GPU加速,但也可以兼容;
  3. 提升样本在线增强的效率,同样可以通过NVIDIA-DALI库,实现GPU加速;
  4. 在网络结构确定的情况下,提速主要有两种方式,并且可以同时采用:
    • 采用Data Parallel的多卡并行训练
    • 采用amp自动混合精度训练

2. 实验配置

2.1 服务器

服务器为4TITAN RTX,进行实验时停止了其他高资源消耗的进程。

2.2 基本配置

  • Dataset:ImageNet
  • Model:MobilenetV2
  • Augmentation:RandomCrop,RandomFlip,Resize,Normalization
  • 每个进程的batch_size:256
  • 每个进程的Dataloadernum_threads:8

3. 具体实现中的注意点

3.1 关于Dataloader

在使用DALI库构建Dataloader时,建议采用ops.ExternalSource的方式来加载数据,这样可以实现比较高的自由度。

示例代码中只实现了分类任务的dataloader,但按照这种方式构建,很容易实现其他如检测/分割任务的dataloader。

只要把数据来源按照迭代器来实现,就可以套用到ops.ExternalSource这一套框架下。

参见src/datasets/cls_dataset_dali.py中的ClsInputIterator

3.2 关于Loss

在训练过程中,每个进程分别计算各自的loss,通过内部同步机制去同步loss信息。但是在训练中需要监控过程,此时需要计算所有loss的均值。

参见src/train/logger.py中关于reduce_tensor的计算方式。

3.3 关于多进程参数的选取

在训练过程中,实验用的服务器,CPU共32核心,4卡并行,因此每个进程的Dataloader,设定的num_threads为8,测试下来效率最高。
如果num_gpusnum_threads < CPU核心数,不能充分利用CPU资源
如果num_gpus
num_threads > CPU核心数,速度反而也会有所下降。

4. 训练速度实测结果

4.1 未开启amp时的GPU占用

轻量化模型训练加速的思考(Pytorch实现)

4.2 开启amp后的GPU占用

轻量化模型训练加速的思考(Pytorch实现)

4.3 CPU占用情况

开启/关闭amp对于CPU的影响不大,基本看不出区别

轻量化模型训练加速的思考(Pytorch实现)

4.4 综合训练速度

4卡并行,BS为256,训练集约120W图片。训练速度为:

  • 未开启amp:约 2.4 iters/s(2458 帧/s),每个epoch训练时间不到 9 min;
  • 开启amp:约 3.8 iters/s(3891 帧/s),每个epoch训练时间不到 6 min;

5. 一些总结

通过综合采用各种训练加速手段,基本可以做到充分利用多显卡服务器的GPUCPU资源,不会造成硬件资源的浪费;

  1. 通过Nvidia-Dali模块的合理配置,可以显著提升数据加载和在线增强阶段的效率,特别是在训练一些轻量化模型时,往往瓶颈不在于GPU的计算速度,而在于CPU等其他部件的负载;
  2. 通过DistributedDataParallel模块的合理配置,可以实现多卡的负载均衡,不论是显存占用还是GPU利用率,都能够达到平衡,不会有其中1张卡变成效率瓶颈;
  3. 通过torch.cuda.amp模块的合理配置,可以进一步降低显存占用,从而可以设置更大的batch_size,提高模型收敛速度;
  4. torch.cuda.amp模块还可以显著降低网络前向推理时间,从而进一步提高训练效率。

综合应用如上所述的手段,基本上可以实现显卡数量和训练效率之间的线性增长关系
不会发生卡多了,但是单卡的效率却大大下降的现象。

6. 一些意外

原以为本篇到此就该结束了,但又遇到了新的问题。

当训练执行一段时间后,由于整个系统长时间处于高负载的状态,显卡温度飙升,触发了显卡的保护机制,自动降频了,GPU利用率直接降到了原来的一半左右。

之前显卡运行效率低的时候,散热不良的问题还没有显露出来,一旦长时间高负荷运转,多卡密集排布和风冷散热的不足就暴露出来了。

下一步是要折腾水冷散热了么?


喜欢 (0)