点击上方“CVer”,选择加"星标"或“置顶”
重磅干货,第一时间送达
作者:风车车
https://zhuanlan.zhihu.com/p/77633542
本文已由作者授权,未经允许,不得二次转载
前段时间训练了不少模型,发现并不是大力出奇迹,显卡越多越好,有时候1张v100和2张v100可能没有什么区别,后来发现瓶颈在其他地方,写篇文章来总结一下自己用过的一些小trick,最后的效果就是在cifar上面跑vgg的时间从一天缩到了一个小时,imagenet上跑mobilenet模型只需要2分钟每个epoch。(文章末尾有代码啦)
先说下跑cifar的时候,如果只是用torchvision的dataloader(用最常见的padding/crop/flip做数据增强)会很慢,大概速度是下面这种,600个epoch差不多要一天多才能跑完,并且速度时快时慢很不稳定。
![94a794f4c74952c41a05114d27ce211d.png](https://i-blog.csdnimg.cn/blog_migrate/50f86482e2efcf13517ba284e1511028.jpeg)
我最初以为是IO的原因,于是挂载了一块内存盘,改了一下路径接着用torchvision的dataloader来跑,速度基本没啥变化。。。
![f546f1cdd81e26d450f4eb6a0df5cfec.png](https://i-blog.csdnimg.cn/blog_migrate/1b731aa212fa6db5e00e72f850cac3ac.jpeg)
然后打开资源使用率看了下发现cpu使用率几乎已经满了(只能申请2cpu和一张v100...),但是gpu的使用率非常低,这基本可以确定瓶颈是在cpu的处理速度上了。
![497d3bfb79fa1e9ceae25ed7e2d8a7fc.png](https://i-blog.csdnimg.cn/blog_migrate/6f881382c41c9e24456f3ccdb3b0734c.jpeg)
后来查了一些资料发现nvidia有一个库叫dali可以用gpu来做图像的前处理,从输入,解码到transform的一整套pipeline,看了下常见的操作比如pad/crop之类的还挺全的,并且支持pytorch/caffe/mxnet等各种框架。
![df9103f47f3aea840db1c9cf57cb1122.png](https://i-blog.csdnimg.cn/blog_migrate/b49cce75a974455d177c5b8456d5a63a.jpeg)
可惜在官方文档中没找到cifar的pipeline,于是自己照着imagenet的版本写了个,最初踩了一些坑(为了省事找了个cifar的jpeg版本来解码,发现精度掉得很多还找不到原因,还得从cifar的二进制文件来读取),最后总归是达到了同样的精度,再来看一看速度和资源使用率,总时间直接从一天缩短为一小时,并且gpu使用率高了很多。
![ce5f877fbeafa6bdb5d6309ea7d9ac12.png](https://i-blog.csdnimg.cn/blog_migrate/3f3f8c16806da990ad5a6b1f01b13d26.jpeg)
![7dbef57c947cef9760b9acbce672cf53.png](https://i-blog.csdnimg.cn/blog_migrate/229c8e4c1473a3dc68717ea4b430c705.jpeg)
再说下imagenet的训练加速,最初也是把整个数据集拷到了挂载的内存盘里面(160g大概够用了,从拷贝到解压完成大概10分钟不到),发现同样用torchvision的dataloader训练很不稳定,于是直接照搬了dali官方的dataloader过来,速度也是同样起飞hhhh(找不到当时训练的图片了),然后再配合apex的混合精度和分布式训练,申请4块v100,gpu使用率可以稳定在95以上,8块v100可以稳定在90以上,最后直接上到16张v100和32cpu,大概也能稳定在85左右(看资源使用率发现cpu到顶了,不然估计gpu也能到95以上),16块v100在ImageNet上跑mobilenet只需要2分钟每个epoch。
![766fa6de87816afac0cc07d95f175f18.png](https://i-blog.csdnimg.cn/blog_migrate/a27f6c490587602f8d55377bfd426fd3.jpeg)
写的dataloader放到了github上,我测试的精度跟torchvision的版本差不多,不过速度上会比torchvision快很多,后面有空也会写一些其他常用dataloader的dali版本放上去
https://github.com/tanglang96/DataLoaders_DALI
重磅!CVer学术交流群成立啦
扫码添加CVer助手,可申请加入CVer-目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测和模型剪枝&压缩等群。一定要备注:研究方向+地点+学校/公司+昵称(如目标检测+上海+上交+卡卡)
▲长按加群
▲长按关注我们
麻烦给我一个在看!