我们下载下来的预训练的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"的变量的模型用来做预测。这样虽然不能提升模型的预测速度,但是可以减小模型的内存。