前段時間研究了tflite和量化相關的操作, 經測試量化尤其在具有專門DSP加速的硬體上(比如MTK8183)有著很好的加速效果,大約3X的提升; tensorflow提供了tflite轉化工具toco,使用命令大致如下:

bazel-bin/tensorflow/contrib/lite/toco/toco --input_file=mobilenet_v1_1.0_128_frozen.pb \
  --input_format=TENSORFLOW_GRAPHDEF \
  --output_format=TFLITE \
  --output_file=/tmp/mobilenet_v1_1.0_128.tflite \
  --inference_type=FLOAT \
  --input_data_types=FLOAT \
  --input_arrays=input \
  --output_arrays=MobilenetV1/Predictions/Reshape_1 \
  --input_shapes=1,128,128,3 \
  --logtostderr
 

這裡面遇到的一個坑是,發現用pip直接安裝的tensorflow, 執行這個命令會包error:, 因為option的名字不一樣,那是用python重新封裝的,而從原始碼編譯toco(用官方的bazel工具)是C++的; 上述命令只是轉化為tflite的格式,如果需要量化,只需要改成這樣:

bazel-bin/tensorflow/contrib/lite/toco/toco --input_file=mobilenet_v1_1.0_128_frozen.pb \
  --input_format=TENSORFLOW_GRAPHDEF \
  --output_format=TFLITE \
  --output_file=/tmp/mobilenet_v1_1.0_128.tflite \
  --inference_type=QUANTIZED_UINT8 \
  --input_data_types=QUANTIZED_UINT8 \
  --input_arrays=input \
  --output_arrays=MobilenetV1/Predictions/Reshape_1 \
  --input_shapes=1,128,128,3 \
  --logtostderr
 

對於, 一般的模型, 它會報錯, 原因是沒有提供MinMax的範圍(量化需要這兩個值進行縮放), 對於不同的型別,這兩個值的獲取方式不一樣,對於裡面的weights比較好辦,只要統計一下最大最小值就好了,但是對於graph裡面的其他tensor,因為輸入的不同,所以每層輸出的值也會不同,這裡就需要在訓練的時候去統計每個node的[Min, Max]; 當然你可以用這個語句去手動制定範圍,但是精度可能就會收到很大影響.

tf.quantization.fake_quant_with_min_max_args
而且使用toco的時候需要加上:
--reorder_across_fake_quant=true
 

或者你也可以在使用toco的時候加上:

--default_ranges_min=0 \
--default_ranges_max=1

但是這樣就更加糟糕了,它會給所有沒有[Min,Max]的node,使用這兩個值作為預設值   所以,由此我們知道: 1) 我們需要在訓練的時候統計[Min, Max] 2) 為了量化能有更好的精度,我們希望[Min, Max]最好在一個比較小的範圍內, 所以一些best practise包括: a)使用relu6作為激勵函式; b) 使用batch normalization   Google提供了 quantization-aware的訓練方法,可以參考 https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/quantize 簡單總結下: 1) 在你定義好網路結構之後, 加上

tf.contrib.quantize.create_training_graph(input_graph=g,
                                          quant_delay=2000000)

它會把網路重新改造,具體可以看tensorboar; 這裡面,要注意quant_delay這個引數, 這個引數的意思的在此之前,網路會使用float進行正常的訓練,收斂到一個比較好的水平;之後,會執行一個fake-quant的操作. 改造之前的一個卷積層: 改造之後的一個卷積層: 關鍵的操作就在於這個act_quant節點,統計了最大最小的值,然後執行量化操作,最上面的delay_quant裡面有兩個switch op, 就用根據設定的quant_delay來決定是否使用量化的結果進行forward 2) 按照正常的方式儲存你的checkpoint; 3) 匯出用於eval的網路,這個最好是在另外一個檔案裡面(否則可能會有這個問題 https://github.com/tensorflow/tensorflow/issues/19936),同樣的方式建立網路結構後,呼叫:

tf.contrib.quantize.create_eval_graph(input_graph=g)

然後load checkpoint,最後儲存graph和新的checkpoint.

tf.train.write_graph(sess.graph_def, output_dir, 'graph.pb', as_text=True)
saver = tf.train.Saver(max_to_keep=100)
saver.save(sess, './workspace/'+args.model+'/chk', global_step=1)

4) freeze graph

python3 -m tensorflow.python.tools.freeze_graph \
--input_graph=./graph.pb \
--output_graph=exported_freezed_inference_graph.pb \
--input_checkpoint=./chk-1 \
--output_node_names="your_output_name"

5) 現在就可以使用toco命令匯出量化的模型了;   也許等你執行完以上操作,發現自己的模型裡面仍然會報某些node還是沒有[Min, Max], 根據我的觀察,在某些情況下(具體不明),發現只有添加了BatchNorm的層在呼叫cerate_train_graph的時候才會新增fake_quant; 但是有的時候不需要BatchNorm也可以,很奇怪~   最後,你還可以用這個命令檢視你到處tflite量化模型的op和[Min,Max]的範圍

bazel-bin/tensorflow/contrib/lite/toco/toco \
--input_file=exported_freezed_inference_graph.pb \
--input_format=TENSORFLOW_GRAPHDEF   \
--output_format=GRAPHVIZ_DOT   \
--output_file=exported_freezed_inference_graph_opt.dot   \
--inference_type=QUANTIZED_UINT8 \
--input_data_types=QUANTIZED_UINT8  \
--input_arrays=image   \
--output_arrays=Openpose/MConv_Stage1_L_5_pointwise/Conv2D \
--input_shapes=1,128,128,3   \
--logtostderr \
--default_ranges_min=0 \
--default_ranges_max=1

dot -Tpdf exported_freezed_inference_graph_opt.dot -o exported_freezed_inference_graph_opt.pdf