最近項目里需要一個小型的目標檢測模型，SSD、YOLO等一通模型調參試下來，直接調用TensorFlow object detect API居然效果最好，大廠的產品不得不服啊。使用mobilenet ssd v2模型，配置文件也未修改參數，訓練后的模型不光檢測效果不錯，在CPU上的運行時間也在70ms左右。之后將模型移植到安卓手機上（魅族MX4，老的不是一點點），卡頓明顯；改用同事的華為，在麒麟960上略微流暢了一些，但仍然不能達到實時檢測。而且訓練得到的pb模型居然有19M，實在太大了，于是又探索了一波模型的壓縮和量化。

模型壓縮

說到模型壓縮，最簡單粗暴的方法當然是減少卷積層數。在使用Tensorflow的API之前，我訓練過一個SSD模型，檢測效果不錯，但耗時接近1s。為了提高檢測速度我果斷開始減少卷積層數，并做了不同層數的對比試驗。結果和原始的VGG16骨干相比，要么檢測效果相近，耗時也沒少多少，要么耗時大減，但漏檢率飆升。也就是在這個情況下，我轉投了mobilenet網絡。

所以這次面臨模型壓縮時， 我沒有再嘗試這個選項（當然也有配置文件不支持刪減層數，要刪就要去改slim里的源碼這個原因。我一個前同事是中科院計算機博士，他的格言就是覺得源碼不好就別調用，自己寫；要調用就盡量避免改源碼，因為你肯定沒有源碼寫得好）。這樣看下來，就只能在配置文件的范圍內自由發揮了。

修改配置文件

首先，附上Tensorflow object detection API中支持的各大模型的配置文件地址：

models/research/object_detection/samples/configs at master · tensorflow/models · GitHub

這里面關于mobilenet_ssd_v2的有好幾個：

image

我使用的是最經典的基于COCO數據集訓練的配置文件，也就是第一個。圖里的最后一個也是基于COCO數據集的，不過是有量化的模型，這個文件我在后面也有用到。

打開配置文件，里面主要分成model、train和eval三塊。在調用API訓練自己的數據時，train和eval的數據當然是要修改的：

image

回到model部分，在feature_extractor那里，有一個depth_multiplier，這個參數作為一個因子與網絡中各層的channel數相乘，換言之，depth_multiplier越小，網絡中feature map的channel數越少，模型參數自然也就少了很多。depth_multiplier默認為1，在我的實驗里改成了0.25，試就試一把大的。

訓練模型

之前depth_multiplier為1時， 我訓練是加載了預訓練模型的，模型地址：

models/detection_model_zoo.md at master · tensorflow/models · GitHub

image

從圖中可以看出，mobilenet_v1的預訓練模型中有一種0.75_depth的版本，這就是depth_multiplier取0.75時在COCO數據集上訓練出來的模型。對于mobilenet_v2，只提供了非量化版和量化版（個人覺得應該0.25、0.5、0.75這幾個常用檔都提供一個，難道是官方不建議壓縮太多嗎。。。）

由于沒有對應的預訓練模型，所以可以選擇加載或者不加載模型。

加載模型的話，開始訓練后命令行會打印一大堆XXX is available in checkpoint, but has an incompatible shape with model variable. This variable will not be initialized from the checkpoint. 不過這并不影響訓練，忽略就可以了。

不加載的話，就將配置文件里fine_tune_checkpoint的那兩行注釋掉。

進入到object detection目錄，運行python object_detection/model_main.py --pipeline_config_path=xxxxxxx/ssd_mobilenet_v2_coco.config --model_dir=xxxxxxxx即可

PS：訓練過程中是不會打印訓練信息的，看命令行會以為電腦卡住了。。。直到eval才會打印出信息

PPS：可以通過TensorBoard來監聽訓練過程，判斷訓練是在正常進行還是電腦真的卡住了（這種情況可能是因為batch size和輸入圖片大小太大。默認是24和300*300，但也都可以改）

模型導出

訓練完成之后，還是在object detection目錄下，運行python export_inference_graph.py，必要的參數分別是輸入的ckpt的文件地址，輸出的pb文件的文件夾以及配置文件地址。

在深度壓縮至0.25倍之后， 我的pb模型大小僅為2.2M，效果卓群。當然網絡的縮減會帶來精度的損失，我的AR和AP分別降了2個點和3個點。

模型移植

生成tflite模型

Tensorflow object detection API訓練出的模型，講道理從ckpt轉成tflite只需要兩步：

第一步，將ckpt轉成pb文件，這次使用的是python export_tflite_ssd_graph.py，操作難度不大，會得到tflite_graph.pb和tflite_graph.pbtxt兩個文件；

第二步，將pb轉為tflite文件，我搜到的方法大都是使用bazel編譯tensorflow/contirb/lite/toco下面的toca文件，但我反復嘗試，報了多種錯誤，依舊沒有成功。。。最后我在stackoverflow上搜到了一位小哥的回復，進入tensorflow/contrib/lite/python目錄，運行python tflite_convert.py，參數設置為

--graph_def_file=XXX/tflite_graph.pb 上一步生成的pb文件地址

--output_file=XXX/xxx.tflite 輸出的tflite文件地址

--input_arrays=normalized_input_image_tensor 輸入輸出的數組名稱對于mobilenet ssd是固定的，不用改

--output_arrays='TFLite_Detection_PostProcess','TFLite_Detection_PostProcess:1','TFLite_Detection_PostProcess:2','TFLite_Detection_PostProcess:3'

--input_shape=1,XXX,XXX,3 輸入的圖片大小，需要與配置文件中一致

--allow_custom_ops

驗證tflite模型

在將tflite模型放進手機之前，我在python里加載tflite模型測試了一次，流程類似加載pb模型

第一步，導入模型

interpreter = tf.contrib.lite.Interpreter(model_path="compress_export/detect.tflite")

interpreter.allocate_tensors()

第二步，獲得輸入和輸出的tensor

input_details = interpreter.get_input_details()

output_details = interpreter.get_output_details()

第三步，讀取輸入圖像，feed給輸入tensor

可以采用PIL或cv2將圖像讀入，轉為numpy數組，然后賦值給input_data

input_data = np.array(XXX)`

interpreter.set_tensor(input_details[0]['index'], input_data)

第四步，運行模型

interpreter.invoke()

第五步， 獲得輸出

參考輸入tensor的表示方法，目標檢測的輸出有4個，具體的值可以通過output_details[0]['index']、output_details[1]['index']、output_details[2]['index']、output_details[3]['index']獲得

這里有一個我踩到的坑，驗證tflite模型時，我采用了和加載pb模型完全相同的圖片預處理步驟，輸出的結果完全不同。幾番檢查之后，發現問題出在模型轉換時。運行python tflite_convert.py時，輸入數組的名稱為normalized_input_image_tensor，而我訓練時采用的是未經normalized的數組。所以在模型轉換時，tensorflow內置了對input進行normalized的步驟。因此在調用tflite模型時，同樣需要在圖像預處理中加入這一步。 nomlized的方法為除以128.0再減去1，保證輸入的值在[-1,1)范圍內。

補充：pb轉tflite的python腳本

import matplotlib

matplotlib.use('TkAgg')

import tensorflowas tf


def test():

    # 指定要使用的模型的路徑  包含圖結構，以及參數

    graph_def_file= 'pb文件路徑'

    # 重新定義一個圖

    output_graph_def= tf.GraphDef()

    with tf.gfile.GFile(graph_def_file,'rb')as fid:

        # 將*.pb文件讀入serialized_graph

        serialized_graph= fid.read()

        # 將serialized_graph的內容恢復到圖中

        output_graph_def.ParseFromString(serialized_graph)

        # print(output_graph_def)

        # 將output_graph_def導入當前默認圖中(加載模型)

        tf.import_graph_def(output_graph_def,name='')

    print('模型加載完成')

    # 使用默認圖，此時已經加載了模型

    detection_graph= tf.get_default_graph()

    with tf.Session(graph=detection_graph)as sess:

        '''

        獲取模型中的tensor

        '''

        image_tensor = detection_graph.get_tensor_by_name('input_1:0') #pb模型輸入的名字

        converter= tf.lite.TFLiteConverter.from_frozen_graph(graph_def_file, ['input_1'], ['output_1'], input_shapes={'input_1': [1,112,112,3]}) #pb模型輸入、輸出的名字以及輸入的大小

        # converter.post_training_quantize = True

        tflite_model= converter.convert()

        converter.optimizations= [tf.lite.Optimize.DEFAULT]

        open("tflite文件路徑","wb").write(tflite_model)

if __name__== '__main__':

    test()

參考

[https://blog.csdn.net/qq_26535271/article/details/84930868](https://blog.csdn.net/qq_26535271/article/details/84930868)

[Tensorflow Convert pb file to TFLITE using python - Stack Overflow](https://stackoverflow.com/questions/50632152/tensorflow-convert-pb-file-to-tflite-using-python)