节省显存的重计算技巧也有了Keras版了
By 苏剑林 | 2020-04-29 | 48962位读者 |不少读者最近可能留意到了公众号文章《BERT重计算:用22.5%的训练时间节省5倍的显存开销(附代码)》,里边介绍了一个叫做“重计算”的技巧,简单来说就是用来省显存的方法,让平均训练速度慢一点,但batch_size可以增大好几倍。该技巧首先发布于论文《Training Deep Nets with Sublinear Memory Cost》,其实在2016年就已经提出了,只不过似乎还没有特别流行起来。
探索 #
公众号文章提到该技巧在pytorch和paddlepaddle都有原生实现了,但tensorflow还没有。但事实上从tensorflow 1.8开始,tensorflow就已经自带了该功能了,当时被列入了tf.contrib
这个子库中,而从tensorflow 1.15开始,它就被内置为tensorflow的主函数之一,那就是tf.recompute_grad
。
找到tf.recompute_grad
之后,笔者就琢磨了一下它的用法,经过一番折腾,最终居然真的成功地用起来了,居然成功地让batch_size
从48增加到了144!然而,在继续整理测试的过程中,发现这玩意居然在tensorflow 2.x是失效的...于是再折腾了两天,查找了各种资料并反复调试,最终算是成功地补充了这一缺陷。
最后是笔者自己的开源实现:
该实现已经内置在bert4keras中,使用bert4keras的读者可以升级到最新版本(0.7.5+)来测试该功能。
使用 #
笔者的实现也命名为recompute_grad
,它是一个装饰器,用于自定义Keras层的call
函数,比如
from recompute import recompute_grad
class MyLayer(Layer):
@recompute_grad
def call(self, inputs):
return inputs * 2
对于已经存在的层,可以通过继承的方式来装饰:
from recompute import recompute_grad
class MyDense(Dense):
@recompute_grad
def call(self, inputs):
return super(MyDense, self).call(inputs)
自定义好层之后,在代码中嵌入自定义层,然后在执行代码之前,加入环境变量RECOMPUTE=1
来启用重计算。
注意:不是在总模型里插入了@recompute_grad
,就能达到省内存的目的,而是要在每个层都插入@recompute_grad
才能更好地省显存。简单来说,就是插入的@recompute_grad
越多,就省显存。具体原因请仔细理解重计算的原理。
效果 #
bert4keras 0.7.5+已经内置了重计算,直接传入环境变量RECOMPUTE=1
就会启用重计算,读者可以自行尝试,大概的效果是:
1、在BERT Base版本下,batch_size可以增大为原来的3倍左右;
2、在BERT Large版本下,batch_size可以增大为原来的4倍左右;
3、平均每个样本的训练时间大约增加25%;
4、理论上,层数越多,batch_size可以增大的倍数越大。
环境 #
在下面的环境下测试通过:
tensorflow 1.14 + keras 2.3.1
tensorflow 1.15 + keras 2.3.1
tensorflow 2.0 + keras 2.3.1
tensorflow 2.1 + keras 2.3.1
tensorflow 2.0 + 自带tf.keras
tensorflow 2.1 + 自带tf.keras
确认不支持的环境:
tensorflow 1.x + 自带tf.keras
欢迎报告更多的测试结果。
顺便说一下,强烈建议用keras 2.3.1配合tensorflow 1.x/2.x来跑,强烈不建议使用tensorflow 2.x自带的tf.keras来跑。
参考 #
最后,笔者的实现主要参考自如下两个源码,在此表示感谢。
转载到请包括本文地址:https://kexue.fm/archives/7367
更详细的转载事宜请参考:《科学空间FAQ》
如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。
如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!
如果您需要引用本文,请参考:
苏剑林. (Apr. 29, 2020). 《节省显存的重计算技巧也有了Keras版了 》[Blog post]. Retrieved from https://kexue.fm/archives/7367
@online{kexuefm-7367,
title={节省显存的重计算技巧也有了Keras版了},
author={苏剑林},
year={2020},
month={Apr},
url={\url{https://kexue.fm/archives/7367}},
}
April 30th, 2020
自定义keras层,总感觉不好验证【也不知道自己写的对不对】,苏神是怎么做的
不好验证什么?
May 6th, 2020
试试公式$a^{2} + b^{2} = c^{2}$
May 6th, 2020
请教一下 如果用预训练的模型 比如resnet50 这个重计算怎么用?谢谢!
June 7th, 2020
请教一下,为什么在创建模型(model=Model(inputs,outputs))时报错
AttributeError: 'NoneType' object has no attribute '_inbound_nodes'
呃,别在意这个,这是我从tf.keras切到keras冒出来的东西
不清楚。可能是tf.keras和keras混用了吧
March 9th, 2021
您好,我用的是tensorfolow 1.12的版本,请问重计算在contrib库的什么地方呢?
不知道,我不用tf 1.12
April 20th, 2021
请问苏佬,我搭建模型的时候使用了build_transformer_model函数,我是将Transformer基类中的call方法加上@recompute_grad,但是这样貌似行不通,从32批次增加到64批次报OOM了,请问我应该怎样实现
先认真读读本文。