在深度学习中喂饱GPU

新智元

发布于 2019-08-15 18:20:12

1.7K0

发布于 2019-08-15 18:20:12

文章被收录于专栏：新智元新智元

新智元推荐

来源：知乎专栏

作者：风车车

【新智元导读】深度学习模型训练是不是大力出奇迹，显卡越多越好？非也，没有512张显卡，也可以通过一些小技巧优化模型训练。本文作者分析了他的实践经验。

前段时间训练了不少模型，发现并不是大力出奇迹，显卡越多越好，有时候 1 张 v100 和 2 张 v100 可能没有什么区别，后来发现瓶颈在其他地方，写篇文章来总结一下自己用过的一些小 trick，最后的效果就是在 cifar 上面跑 vgg 的时间从一天缩到了一个小时，imagenet 上跑 mobilenet 模型只需要 2 分钟每个 epoch。（文章末尾有代码啦）

先说下跑 cifar 的时候，如果只是用 torchvision 的 dataloader (用最常见的 padding/crop/flip 做数据增强) 会很慢，大概速度是下面这种，600 个 epoch 差不多要一天多才能跑完，并且速度时快时慢很不稳定。

我最初以为是 IO 的原因，于是挂载了一块内存盘，改了一下路径接着用 torchvision 的 dataloader 来跑，速度基本没啥变化。。。

然后打开资源使用率看了下发现 cpu 使用率几乎已经满了（只能申请 2cpu 和一张 v100...），但是 gpu 的使用率非常低，这基本可以确定瓶颈是在 cpu 的处理速度上了。

后来查了一些资料发现 nvidia 有一个库叫 dali 可以用 gpu 来做图像的前处理，从输入，解码到 transform 的一整套 pipeline，看了下常见的操作比如 pad/crop 之类的还挺全的，并且支持 pytorch/caffe/mxnet 等各种框架。

可惜在官方文档中没找到 cifar 的 pipeline，于是自己照着 imagenet 的版本写了个，最初踩了一些坑（为了省事找了个 cifar 的 jpeg 版本来解码，发现精度掉得很多还找不到原因，还得从 cifar 的二进制文件来读取），最后总归是达到了同样的精度，再来看一看速度和资源使用率，总时间直接从一天缩短为一小时，并且 gpu 使用率高了很多。

再说下 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。