我們下載下來的預訓練的bert模型的大小大概是400M左右,但是我們自己預訓練的bert模型,或者是我們在開源的bert模型上fine-tuning之后的模型的大小大約是1.1G,我們來看看到底是什么原因造成的,首先我們可以通過下一段代碼來輸出我們訓練好的模型的參數變量。
下面這段代碼可以輸出我們下載的官方預訓練模型的參數變量
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模型,代碼和上面一樣
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)
我們同樣截取部分參數
我們可以看到這里混入了不少帶有"adam"的變量,我們來看adam優化算法,在計算一階矩和二階矩時,我們是要保存之前時刻的滑動平均值的,而每個需要通過梯度更新的參數,都要維護這樣一個一階矩和二階矩之前時刻的滑動平均值,也就是對應上面的 "adam_m" (一階矩) 和 “adam_v” (二階矩),因此導致我們自己預訓練的模型的大小大約是官方預訓練模型的大小的3倍。而這些參數變量只有訓練模型的時候有用,在之后預測的時候以及fine-tuning階段都是沒有用的(fine-tuning時我們只是用到了之前預訓練好的模型的參數來作為初始化值,並不會用到優化算法中的中間值),因此我們可以在訓練完或者fine-tuning完bert模型之后,在保存模型時將這些參數去掉,也可以在保存了完整的參數之后,再加載去掉這些參數,然后重新保存,這樣就不需要改動bert的源碼,具體的實現如下:
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"的變量的模型用來做預測。這樣雖然不能提升模型的預測速度,但是可以減小模型的內存。