问题描述
最近课程实验在与同学交流的过程中发现自己的模型训练速度比别人慢很多,而且使用AI Studio的Tesla V100速度与自己笔记本的GTX1050速度差别不大,跑一个ResNet50的epoch要12分钟,一度让我怀疑百度给我提供的是阉割版的显卡。
尤其是训练ResNet和VGG的时候更加发现了这个问题,使用nvidia-smi查看了一下结果如下:
显然GPU利用率为0就很不正常,但是有显存占用说明模型应该是在跑的。后来既然GPU利用不起来,我干脆同时跑多个模型,想充分利用,但是效果还不是很明显,速度依然都很慢。
查了查资料发现可以在Dataloader里加pin_memory=True, num_workers=8,这样能将数据放入显存,并且通过多线程减少数据读入的时间,模型训练速度会快一些。但是我试了试也没有明显的效果,就很奇怪。
当然也试过增大batch size,但是治标不治本。
问题原因
所以到底为什么导致了这个原因呢?
后来我偶然为了减小最后的tensorboard log文件,我把一些用不到的tensorboard的写入比如直方图给删掉了,发现速度瞬间提上来了,然后再减小训练时的输出间隔,速度也有提升。立刻明白了,是因为训练的主要时间都花在了写日志上,文件IO耗时特别多,尤其是我设置的写入间隔还很小,所以GPU计算一瞬间,然后写很久的记录,计算一瞬间,再写很久的记录,最终导致速度特别慢。这也正是为了速度和我笔记本的GTX1050一样,因为大家比的都是写文件的速度。
最后的效果还是很好的,不愧是Tesla V100。ResNet50从最初跑一个epoch12分钟提升到四十多秒,如果删掉日志写入还能更快。而VGG16也只要二十多秒,训练100个epoch分分钟的事情。
另外这么快的速度也不全是显卡,跟前面提到的pin_memory也有关系,以及使用了PyTorch的半精度来提高速度,精度也没有受到很大影响。
这个问题解决后我也明白为啥之前的AlexNet跑一个epoch还要三四分钟,五六分钟,时间都花在写日志上了。以及之前跑的一个比较大的模型,也有了提速的方法,之前需要50分钟,现在估计五六分钟。可以很方便的做各种对比实验与测试,效率提高的不是一点点。
总结
最后总结一下,有的时候模型训练慢并不是因为显卡不行或者模型太大,而是在跑模型过程中有一些其他的操作导致速度很慢,尤其是文件的IO操作,这会导致GPU得不到连续性使用,整体速度特别慢。
模型提速技巧
- 减少日志IO操作频率
- 使用pin_memory和num_workers
- 使用半精度训练
- 更好的显卡,更轻的模型
另外也可以通过增大batch size提高epoch速度,但是收敛速度也会变慢,需要再适当升高学习率