tensorflow環境下實現bert_base量化，完成bert輕量級

環境：

windows 10

python 3.5

GTX 1660Ti

tensorflow-gpu 1.13.1

numpy  1.18.1　　

 

1. 首先下載google開源的預訓練好的model。我本次用的是 BERT-Base, Uncased(第一個)

　　BERT-Base, Uncased: 12-layer, 768-hidden, 12-heads, 110M parameters
　　BERT-Large, Uncased: 24-layer, 1024-hidden, 16-heads, 340M parameters
　　BERT-Base, Cased: 12-layer, 768-hidden, 12-heads , 110M parameters
2. 參考https://www.zybuluo.com/Team/note/1632532  (  https://zhuanlan.zhihu.com/p/91024786?utm_source=wechat_session&utm_medium=social&utm_oi=1035849572991401984)第四部分。也就是下圖：

　　

 3. 打開nvidia的github官方，對其中的convert部分進行參數修改。

　　https://github.com/NVIDIA/DeepLearningExamples/blob/master/FasterTransformer/v2/sample/tensorflow_bert/ckpt_type_convert.py

　　其中代碼詳情如下：(實現的功能就是將 FP32 convert to  FP16)

import tensorflow as tf
import numpy as np
from tensorflow.contrib.framework.python.framework import checkpoint_utils
from tensorflow.python.ops import io_ops
from tensorflow.python.training.saver import BaseSaverBuilder


def checkpoint_dtype_cast(in_checkpoint_file, out_checkpoint_file):
    var_list = checkpoint_utils.list_variables(tf.flags.FLAGS.init_checkpoint)

    def init_graph():
        for name, shape in var_list:
            var = checkpoint_utils.load_variable(tf.flags.FLAGS.init_checkpoint, name)
            recon_dtype = tf.float16 if var.dtype == np.float32 else var.dtype
            tf.get_variable(name, shape=shape, dtype=recon_dtype)

    init_graph()
    saver = tf.train.Saver(builder=CastFromFloat32SaverBuilder())
    with tf.Session() as sess:
        saver.restore(sess, in_checkpoint_file)
        saver.save(sess, 'tmp.ckpt')

    tf.reset_default_graph()

    init_graph()
    saver = tf.train.Saver()
    with tf.Session() as sess:
        saver.restore(sess, 'tmp.ckpt')
        saver.save(sess, out_checkpoint_file)


class CastFromFloat32SaverBuilder(BaseSaverBuilder):
    # Based on tensorflow.python.training.saver.BulkSaverBuilder.bulk_restore
    def bulk_restore(self, filename_tensor, saveables, preferred_shard,
                     restore_sequentially):
        restore_specs = []
        for saveable in saveables:
            for spec in saveable.specs:
                restore_specs.append((spec.name, spec.slice_spec, spec.dtype))
        names, slices, dtypes = zip(*restore_specs)
        restore_dtypes = [tf.float32 if dtype.base_dtype==tf.float16 else dtype for dtype in dtypes]
        # print info
        for i in range(len(restore_specs)):
            print(names[i], 'from', restore_dtypes[i], 'to', dtypes[i].base_dtype)
        with tf.device("cpu:0"):
            restored = io_ops.restore_v2(
                filename_tensor, names, slices, restore_dtypes)
            return [tf.cast(r, dt.base_dtype) for r, dt in zip(restored, dtypes)]


if __name__ == '__main__':
    tf.flags.DEFINE_string("fp16_checkpoint", "mrpc_output/fp16_model.ckpt", "fp16 checkpoint file")
    tf.flags.DEFINE_string("init_checkpoint", "bert_base/bert_model.ckpt", "initial checkpoint file")
    checkpoint_dtype_cast(tf.flags.FLAGS.init_checkpoint, tf.flags.FLAGS.fp16_checkpoint)

　　其中,main函數的兩個參數可以修改。第一個是你convert后的模型最終要輸出的地方，第二個是你下載的google的模型的地址.

 

很多讀者看到這里就覺得多此一舉，為什么不知將通過tensorflow的官方工具，只需要幾行代碼就可以實現float32->float16(不知道的小伙伴可以看這里)，但是需要注意的是，使用TFLite轉換得到的量化模型是tflite結構，意味着只能在tflite中運行（大部分場景為移動端）具體可以參考這里。

 4. 

　　通過量化后的bert模型，我們就可以進行測試性能了。身邊正好有一個錯別字中ppl的計算model，所以把bert量化后，直接可以進行性能測試。

　　(注：由於restore模型時，是通過先加載運算圖，再加載圖中的變量參數等信息，有很多錯誤。因此我們需要縷一遍代碼，將其中的variable的dtype修改為float16，否則出現類型不一致等錯誤。)

　　之后通過bert作為語言模型計算每個句子的ppl的性能和時間作為評測標准，進行了模型輕量級前后的比較：

　　轉換前float32計算一個句子的ppl和時間和顯存占用：

 

 　　轉化后float16:

　

5.

　　最后我們得出結論：量化后的模型相對於原模型精度會有些許損失，但是顯存占用減少了很多。

 

 

參考；   https://zhuanlan.zhihu.com/p/113734249