如果自己有一批数据,想要使用预训练模型再进行领域数据的适配,可以使用huggingface中的run_mlm.py进行训练。
最近在训练的时候遇到一个问题:训练的时候GPU占用上升导致out of memory。
详细问题描述:pytorch 使用 transoformers(version=4.8.2)中 run_mlm.py 进行垂直领域数据适配预训练的时候;GPU显存占用量慢慢上升,但是并不是一直在升,隔几百个step之后会有一个量(1G)的上升并稳定一段时间后再次上升; 直至 out of memory。
代码的报错位置是 :
网上查了下,看到了这个解答:
和报错位置的代码差不多,但是这个问题应该是pytorch之前版本的问题了,新的版本应该已经修复这个问题了;试了一下解决方法,果然还是有上述问题。
接着就看到了这个博客,对局部的损失loss进行了删除并回收。
加入了这两行语句,简单暴力,再看gpu的使用会在两个值之间跳动,应该是代码生效了。
总结:原因应该是,损失累加的时候有很多局部的损失没有删除;因为训练循环没有结束,所以一直占用显存。