Bert系列模型参数优化降低大小

我们下载下来的预训练的bert-base模型的大小大概是394M左右,但我们在自己数据集上经过fine-tuning后的bert-bae模型大小大约是1.2G, 整整是Bert-base模型的3倍,让我们来看看到底是什么原因造成的,首先我们可以通过下一段代码来输出我们训练好的模型和官方提供的Bert-base模型的参数变量。

1:官方提供的Bert-base模型参数信息如下:
from tensorflow.python import pywrap_tensorflow
from tensorflow.contrib.slim import get_variables_to_restore
# display the bert hyperparameter
def display_bert_base_hyperparameter():
    bert_model_path="/chinese_L-12_H-768_A-12/"
    model_reader = pywrap_tensorflow.NewCheckpointReader(bert_model_path + "bert_model.ckpt")
    var_dict = model_reader.get_variable_to_shape_map()
    print(len(var_dict))
    for key,value in var_dict.items():
        print(str(key) + ":" + str(value))
    print("=*="*10)

Bert-base部分参数如下所示在这里插入图片描述

2:我们自己fine-tuning后的Bert-base模型参数信息如下:

在这里插入图片描述
通过对比分析发现混入了不少带有"adam"的变量,我们来看adam优化算法,在计算一阶矩和二阶矩时,我们是要保存之前时刻的滑动平均值的,而每个需要通过梯度更新的参数,都要维护这样一个一阶矩和二阶矩之前时刻的滑动平均值,也就是对应上面的 “adam_m” (一阶矩) 和 “adam_v” (二阶矩),因此导致我们自己预训练的模型的大小大约是官方预训练模型的大小的3倍。而这些参数变量只有训练模型的时候有用,在之后预测的时候以及fine-tuning阶段都是没有用的(fine-tuning时我们只是用到了之前预训练好的模型的参数来作为初始化值,并不会用到优化算法中的中间值),因此我们可以在训练完或者fine-tuning完bert模型之后,在保存模型时将这些参数去掉,也可以在保存了完整的参数之后,再加载去掉这些参数,然后重新保存,这样就不需要改动bert的源码,具体的实现如下:

def slim_bert_model_size():
    bert_fintune_model_path = "/fine-tuning-model/"
    graph = tf.Graph()
    with graph.as_default():
        sess = tf.Session()
        checkpoint_file = tf.train.latest_checkpoint(bert_fintune_model_path)
        saver = tf.train.import_meta_graph("{}.meta".format(checkpoint_file))
        saver.restore(sess, checkpoint_file)

        variables = get_variables_to_restore()
        other_vars = [variable for variable in variables if not re.search("adam", variable.name)]
        var_saver = tf.train.Saver(other_vars)
        var_saver.save(sess, bert_fintune_model_path + "model.ckpt-slim")

之后就可以直接加载这个去掉带"adam"的变量的模型用来做预测。这样虽然不能提升模型的预测速度,但是可以减小模型的内存。

  • 1
    点赞
  • 5
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值