Bert層數剪枝

模型精簡的流程如下：pretrian model -> retrain with new data(fine tuning) -> pruning -> retrain -> model

對bert進行層數剪枝，保留第一層和第十二層參數，再用領域數據微調。代碼如下：

"""
    test
"""
import tensorflow as tf
import os

sess = tf.Session()
last_name = 'bert_model.ckpt'
model_path = 'bert_model/chinese_L-12_H-768_A-12'
imported_meta = tf.train.import_meta_graph(os.path.join(model_path, last_name + '.meta'))
imported_meta.restore(sess, os.path.join(model_path, last_name))
init_op = tf.local_variables_initializer()
sess.run(init_op)

bert_dict = {}
# 獲取待保存的層數節點
for var in tf.global_variables():
    # print(var)
    # 提取第0層和第11層和其它的參數，其余1-10層去掉，存儲變量名的數值
    if var.name.startswith('bert/encoder/layer_') and not var.name.startswith(
            'bert/encoder/layer_0') and not var.name.startswith('bert/encoder/layer_11'):
        pass
    else:
        bert_dict[var.name] = sess.run(var).tolist()

# print('bert_dict:{}'.format(bert_dict))
# 真是保存的變量信息
need_vars = []
for var in tf.global_variables():
    if var.name.startswith('bert/encoder/layer_') and not var.name.startswith(
            'bert/encoder/layer_0/') and not var.name.startswith('bert/encoder/layer_1/'):
        pass
    elif var.name.startswith('bert/encoder/layer_1/'):
        # 尋找11層的var name，將11層的參數給第一層使用
        new_name = var.name.replace("bert/encoder/layer_1", "bert/encoder/layer_11")
        op = tf.assign(var, bert_dict[new_name])
        sess.run(op)
        need_vars.append(var)
        print(var)
    else:
        need_vars.append(var)
        print('####',var)

# 保存model
saver = tf.train.Saver(need_vars)
saver.save(sess, os.path.join('bert_model/chinese_L-12_H-768_A-12_pruning', 'bert_pruning_2_layer.ckpt'))

 要修改對應的配置文件參數：

 

效果總結

在bert_base版本二分類模型的F1值達到97%，經過該方法裁剪后F1達到93.99%，損失在3個點左右，符合預期，還是可以投入工程使用的