tensorflow的freeze graph及inference graph_transforms

tensorflow模型優化

一般利用tensorflow訓練出來的模型包含4個文件：

• check_point文件，指定模型存放路徑的文件
• .data文件，存放模型的權重文件
• .index文件，存放網絡節點的索引
• .meta文件，存放網絡圖結構

當我們在自己的電腦上實驗時，固然可以先加載meta文件獲得網絡圖結構，再加載.data文件加載權重，然後進行推理(inference)。但是在生產環境下，這樣做就有些麻煩，況且有些模型還需要放在移動端，這就必須要優化模型。

1. 固化模型（將權重和圖整合在一起，並去除與inference無關的節點）

源碼位置：tensorflow/python/tools/freeze_graph.py
首先利用bazel編譯tensorflow，bazel build --config=opt --config=cuda //tensorflow/tools/pip_package:build_pip_package
然後再編譯freeze_graph，bazel build tensorflow/python/tools:freeze_graph
用法：

bazel-bin/tensorflow/python/tools/freeze_graph \
--input_graph=some_graph_def.pb \    # 注意：這裏的pb文件是用tf.train.write_graph方法保存的圖結構
--input_checkpoint=model.ckpt \
--output_graph=frozen_graph.pb
--output_node_names=output_node

所以我們在訓練的時候需要用saver.save保存ckpt，然後再用tf.train.write_graph保存圖結構，代碼如下：

with tf.Session() as sess:
    saver = tf.train.Saver()
    saver.save(session, "model.ckpt")
    tf.train.write_graph(session.graph_def, '', 'graph.pb')

如果保存的圖結構是.meta文件怎麼辦呢，也就是說用Saver.save方法和checkpoint一起生成的元模型文件,freeze_graph.py不適用，但可以改源碼，然後重新用bazel編譯。
修改freeze_graph.py裏的_parse_input_graph_proto()函數

# 在freeze_graph.py最前面添加 from tensorflow.python.framework import meta_graph
# 將
input_graph_def = graph_pb2.GraphDef()
mode = "rb" if input_binary else "r"
with gfile.FastGFile(input_graph, mode) as f:
    if input_binary:
      input_graph_def.ParseFromString(f.read())
    else:
      text_format.Merge(f.read(), input_graph_def)
 改爲：
 input_graph_def = meta_graph.read_meta_graph_file(input_graph).graph_def

2.利用tensorflow自帶優化工具實現優化(graph_transforms)

源碼位置：tensorflow/tools/graph_transforms/
github地址
首先編譯graph_transforms
bazel build tensorflow/tools/graph_transforms:transform_graph
查看graph_transforms可知它分爲四個大類：

• Optimizing for Deployment
• Fixing Missing Kernel Errors on Mobile
• Shrinking File Size
• Eight-bit Calculations

2.1 Optimizing for Deployment

當我們在訓練模型後，希望將其部署到服務器或移動設備上，並且希望它儘可能快地運行。該方法刪除了推理過程中沒有調用的所有節點，將始終不變的表達式縮減爲單個節點，並通過對卷積的權值進行預乘，優化了batchnorm過程中使用的一些乘法操作。用法如下：

bazel-bin/tensorflow/tools/graph_transforms/transform_graph \
--in_graph=tensorflow_inception_graph.pb \     #已經固化後的訓練模型
--out_graph=optimized_inception_graph.pb \
--inputs='input,phase_train' \
--outputs='softmax' \
--transforms='
  strip_unused_nodes(type=float, shape="1,299,299,3")
  remove_nodes(op=Identity, op=CheckNumerics)
  fold_constants(ignore_errors=true)
  fold_batch_norms
  fold_old_batch_norms'

2.2 Fixing Missing Kernel Errors on Mobile

由於TensorFlow在Mobile中使用的時候，默認情況下是隻能推理預測，可是在build so文件、jar文件或者.a文件的時候，依賴文件是寫入在tensorflow / contrib / makefile / tf_op_files.txt中的，裏面還包含了一些training相關的ops，這些可能會導致我們在加載PB文件的時候報錯(No OpKernel was registered to support Op)，這時候我們可以通過這個腳本來修復。

bazel-bin/tensorflow/tools/graph_transforms/transform_graph \
--in_graph=tensorflow_inception_graph.pb \
--out_graph=optimized_inception_graph.pb \
--inputs='input,phase_train' \
--outputs='softmax' \
--transforms='
  strip_unused_nodes(type=float, shape="1,299,299,3")
  fold_constants(ignore_errors=true)
  fold_batch_norms
  fold_old_batch_norms'

2.3 Shrinking File Size

如果我們需要將模型部署爲移動應用程序的一部分，那麼我們就需要減小模型的大小啦。對於大多數TensorFlow模型，文件大小的最大貢獻者是傳遞給卷積和全連接層的權重，所以我們要減小模型大小，就得改變權重的存儲方式；默認情況下，權值存儲爲32位浮點值，我們可以用四捨五入的方式存儲權重。

bazel-bin/tensorflow/tools/graph_transforms/transform_graph \
--in_graph=tensorflow_inception_graph.pb \
--out_graph=optimized_inception_graph.pb \
--inputs='input,phase_train' \
--outputs='softmax' \
--transforms='
  strip_unused_nodes(type=float, shape="1,299,299,3")
  fold_constants(ignore_errors=true)
  fold_batch_norms
  fold_old_batch_norms
  round_weights(num_steps=256)'

我們也可以直接把權重量化爲8位來存儲，與round_weights相比，這種方法的缺點是插入了額外的解壓縮操作，將8位值轉換回浮點值，但是TensorFlow運行時中的優化應該確保緩存了這些結果，因此您不應該看到圖形運行得更慢。優化後的模型大小應該約爲原來模型的1/4。

bazel-bin/tensorflow/tools/graph_transforms/transform_graph \
--in_graph=tensorflow_inception_graph.pb \
--out_graph=optimized_inception_graph.pb \
--inputs='input,phase_train' \
--outputs='softmax' \
--transforms='
  strip_unused_nodes(type=float, shape="1,299,299,3")
  fold_constants(ignore_errors=true)
  fold_batch_norms
  fold_old_batch_norms
  quantize_weights'

2.4 Eight-bit Calculations

將推理的計算過程轉換爲8bit定點運算(還在實驗開發階段)

bazel-bin/tensorflow/tools/graph_transforms/transform_graph \
--in_graph=tensorflow_inception_graph.pb \
--out_graph=optimized_inception_graph.pb \
--inputs='input,phase_train' \
--outputs='softmax' \
--transforms='
  add_default_attributes
  strip_unused_nodes(type=float, shape="1,299,299,3")
  remove_nodes(op=Identity, op=CheckNumerics)
  fold_constants(ignore_errors=true)
  fold_batch_norms
  fold_old_batch_norms
  quantize_weights
  quantize_nodes
  strip_unused_nodes
  sort_by_execution_order'

該過程轉換圖中所有具有8位量化的操作，其餘的操作保留在浮點數中。它只支持部分ops，而且在許多平臺上，量化代碼實際上可能比浮點代碼慢，但是當所有條件都合適時，這是一種大幅度提高性能的方法。我管這個叫後量化推理操作，後面我會講到在訓練中加入量化節點的方式完成量化(Quantization-aware training，也叫fake quantization)，見我的另一篇文章。