tensorflow如何使用freeze_graph.py將ckpt轉(zhuǎn)為pb文件

這篇文章主要為大家展示了tensorflow如何使用freeze_graph.py將ckpt轉(zhuǎn)為pb文件，內(nèi)容簡而易懂，希望大家可以學(xué)習(xí)一下，學(xué)習(xí)完之后肯定會(huì)有收獲的，下面讓小編帶大家一起來看看吧。

廢話少說直接上代碼樣例如下

import tensorflow as tf
import os
from tensorflow.python.tools import freeze_graph
# 本來這個(gè)model本無需解釋太多，但是這么多人不能耐下心來看，那么我簡單的說一下吧
# network是你們自己定義的模型結(jié)構(gòu)而已
# ps:
# def network(input)：
# return tf.layers.max_pooling2d(input, 2, 2)
from model import network


os.environ['CUDA_VISIBLE_DEVICES']='2' #設(shè)置GPU


model_path = "path to /model.ckpt-0000" #設(shè)置model的路徑，因新版tensorflow會(huì)生成三個(gè)文件，只需寫到數(shù)字前


def main():

 tf.reset_default_graph()

 input_node = tf.placeholder(tf.float32, shape=(228, 304, 3)) #這個(gè)是你送入網(wǎng)絡(luò)的圖片大小，如果你是其他的大小自行修改
 input_node = tf.expand_dims(input_node, 0)
 flow = network(input_node)
 flow = tf.cast(flow, tf.uint8, 'out') #設(shè)置輸出類型以及輸出的接口名字，為了之后的調(diào)用pb的時(shí)候使用

 saver = tf.train.Saver()
 with tf.Session() as sess:

  saver.restore(sess, model_path)

  #保存圖
  tf.train.write_graph(sess.graph_def, 'output_model/pb_model', 'model.pb')
  #把圖和參數(shù)結(jié)構(gòu)一起
  freeze_graph.freeze_graph('output_model/pb_model/model.pb', '', False, model_path, 'out','save/restore_all', 'save/Const:0', 'output_model/pb_model/frozen_model.pb', False, "")

 print("done")

if __name__ == '__main__':
 main()

這節(jié)是關(guān)于tensorflow的Freezing，字面意思是冷凍，可理解為整合合并；整合什么呢，就是將模型文件和權(quán)重文件整合合并為一個(gè)文件，主要用途是便于發(fā)布。

官方解釋可參考：https://www.tensorflow.org/extend/tool_developers/#freezing 

這里我按我的理解翻譯下，不對(duì)的地方請(qǐng)指正：
有一點(diǎn)令我們?yōu)楸容^困惑的是，tensorflow在訓(xùn)練過程中，通常不會(huì)將權(quán)重?cái)?shù)據(jù)保存的格式文件里（這里我理解是模型文件），反而是分開保存在一個(gè)叫checkpoint的檢查點(diǎn)文件里，當(dāng)初始化時(shí)，再通過模型文件里的變量Op節(jié)點(diǎn)來從checkoupoint文件讀取數(shù)據(jù)并初始化變量。這種模型和權(quán)重?cái)?shù)據(jù)分開保存的情況，使得發(fā)布產(chǎn)品時(shí)不是那么方便，所以便有了freeze_graph.py腳本文件用來將這兩文件整合合并成一個(gè)文件。
freeze_graph.py是怎么做的呢？首行它先加載模型文件，再從checkpoint文件讀取權(quán)重?cái)?shù)據(jù)初始化到模型里的權(quán)重變量，再將權(quán)重變量轉(zhuǎn)換成權(quán)重 常量 （因?yàn)?常量 能隨模型一起保存在同一個(gè)文件里），然后再通過指定的輸出節(jié)點(diǎn)將沒用于輸出推理的Op節(jié)點(diǎn)從圖中剝離掉，再重新保存到指定的文件里（用write_graphdef或Saver）

文件目錄：tensorflow/python/tools/free_graph.py
測試文件：tensorflow/python/tools/free_graph_test.py 這個(gè)測試文件很有學(xué)習(xí)價(jià)值

參數(shù)：

總共有11個(gè)參數(shù)，一個(gè)個(gè)介紹下(必選： 表示必須有值；可選： 表示可以為空)：
1、input_graph：（必選）模型文件，可以是二進(jìn)制的pb文件，或文本的meta文件，用input_binary來指定區(qū)分（見下面說明）
2、input_saver：（可選）Saver解析器。保存模型和權(quán)限時(shí)，Saver也可以自身序列化保存，以便在加載時(shí)應(yīng)用合適的版本。主要用于版本不兼容時(shí)使用?？梢詾榭?，為空時(shí)用當(dāng)前版本的Saver。
3、input_binary：（可選）配合input_graph用，為true時(shí)，input_graph為二進(jìn)制，為false時(shí)，input_graph為文件。默認(rèn)False
4、input_checkpoint：（必選）檢查點(diǎn)數(shù)據(jù)文件。訓(xùn)練時(shí)，給Saver用于保存權(quán)重、偏置等變量值。這時(shí)用于模型恢復(fù)變量值。
5、output_node_names：（必選）輸出節(jié)點(diǎn)的名字，有多個(gè)時(shí)用逗號(hào)分開。用于指定輸出節(jié)點(diǎn)，將沒有在輸出線上的其它節(jié)點(diǎn)剔除。
6、restore_op_name：（可選）從模型恢復(fù)節(jié)點(diǎn)的名字。升級(jí)版中已棄用。默認(rèn)：save/restore_all
7、filename_tensor_name：（可選）已棄用。默認(rèn)：save/Const:0
8、output_graph：（必選）用來保存整合后的模型輸出文件。
9、clear_devices：（可選），默認(rèn)True。指定是否清除訓(xùn)練時(shí)節(jié)點(diǎn)指定的運(yùn)算設(shè)備（如cpu、gpu、tpu。cpu是默認(rèn)）
10、initializer_nodes：（可選）默認(rèn)空。權(quán)限加載后，可通過此參數(shù)來指定需要初始化的節(jié)點(diǎn)，用逗號(hào)分隔多個(gè)節(jié)點(diǎn)名字。
11、variable_names_blacklist：（可先）默認(rèn)空。變量黑名單，用于指定不用恢復(fù)值的變量，用逗號(hào)分隔多個(gè)變量名字。

用法：

例：python tensorflow/python/tools/free_graph.py \
–input_graph=some_graph_def.pb \ 注意：這里的pb文件是用tf.train.write_graph方法保存的
–input_checkpoint=model.ckpt.1001 \ 注意：這里若是r12以上的版本，只需給.data-00000….前面的文件名，如：model.ckpt.1001.data-00000-of-00001，只需寫model.ckpt.1001
–output_graph=/tmp/frozen_graph.pb
–output_node_names=softmax

另外，如果模型文件是.meta格式的，也就是說用saver.Save方法和checkpoint一起生成的元模型文件，free_graph.py不適用，但可以改造下：
1、copy free_graph.py為free_graph_meta.py
2、修改free_graph.py，導(dǎo)入meta_graph：from tensorflow.python.framework import meta_graph
3、將91行到97行換成：input_graph_def = meta_graph.read_meta_graph_file(input_graph).graph_def

這樣改即可加載meta文件

以上就是關(guān)于tensorflow如何使用freeze_graph.py將ckpt轉(zhuǎn)為pb文件的內(nèi)容，如果你們有學(xué)習(xí)到知識(shí)或者技能，可以把它分享出去讓更多的人看到。