咱们下载下来的预训练的bert模型的大小大概是400M左右,可是咱们本身预训练的bert模型,或者是咱们在开源的bert模型上fine-tuning以后的模型的大小大约是1.1G,咱们来看看究竟是什么缘由形成的,首先咱们能够经过下一段代码来输出咱们训练好的模型的参数变量。python
下面这段代码能够输出咱们下载的官方预训练模型的参数变量算法
import tensorflow as tf from tensorflow.python import pywrap_tensorflow model_reader = pywrap_tensorflow.NewCheckpointReader("chinese_L-12_H-768_A-12/bert_model.ckpt") var_dict = model_reader.get_variable_to_shape_map() for key in var_dict: print(key)
咱们截取了部分参数以下:优化
如今换成咱们本身预训练的bert模型,代码和上面同样spa
from tensorflow.python import pywrap_tensorflow model_reader = pywrap_tensorflow.NewCheckpointReader("H_12_768_L12_vocab5/model.ckpt-1500000") var_dict = model_reader.get_variable_to_shape_map() for key in var_dict: print(key)
咱们一样截取部分参数3d
咱们能够看到这里混入了很多带有"adam"的变量,咱们来看adam优化算法,在计算一阶矩和二阶矩时,咱们是要保存以前时刻的滑动平均值的,而每一个须要经过梯度更新的参数,都要维护这样一个一阶矩和二阶矩以前时刻的滑动平均值,也就是对应上面的 "adam_m" (一阶矩) 和 “adam_v” (二阶矩),所以致使咱们本身预训练的模型的大小大约是官方预训练模型的大小的3倍。而这些参数变量只有训练模型的时候有用,在以后预测的时候以及fine-tuning阶段都是没有用的(fine-tuning时咱们只是用到了以前预训练好的模型的参数来做为初始化值,并不会用到优化算法中的中间值),所以咱们能够在训练完或者fine-tuning完bert模型以后,在保存模型时将这些参数去掉,也能够在保存了完整的参数以后,再加载去掉这些参数,而后从新保存,这样就不须要改动bert的源码,具体的实现以下:rest
import re import tensorflow as tf from tensorflow.contrib.slim import get_variables_to_restore # 将bert中和adam相关的参数的值去掉,较小模型的内存 graph = tf.Graph() with graph.as_default(): sess = tf.Session() checkpoint_file = tf.train.latest_checkpoint("H_12_768_L12_vocab5/") 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, "light_bert/model.ckpt")
以后就能够直接加载这个去掉带"adam"的变量的模型用来作预测。这样虽然不能提高模型的预测速度,可是能够减少模型的内存。code