tensorflow實戰筆記(19)----使用freeze_graph.py將ckpt轉為pb文件


一、作用:

https://blog.csdn.net/yjl9122/article/details/78341689

這節是關於tensorflow的Freezing,字面意思是冷凍,可理解為整合合並;整合什么呢,就是將模型文件和權重文件整合合並為一個文件,主要用途是便於發布。
官方解釋可參考:https://www.tensorflow.org/extend/tool_developers/#freezing
這里我按我的理解翻譯下,不對的地方請指正:
有一點令我們為比較困惑的是,tensorflow在訓練過程中,通常不會將權重數據保存的格式文件里(這里我理解是模型文件),反而是分開保存在一個叫checkpoint的檢查點文件里,當初始化時,再通過模型文件里的變量Op節點來從checkoupoint文件讀取數據並初始化變量。這種模型和權重數據分開保存的情況,使得發布產品時不是那么方便,所以便有了freeze_graph.py腳本文件用來將這兩文件整合合並成一個文件。
freeze_graph.py是怎么做的呢?首行它先加載模型文件,再從checkpoint文件讀取權重數據初始化到模型里的權重變量,再將權重變量轉換成權重 常量 (因為 常量 能隨模型一起保存在同一個文件里),然后再通過指定的輸出節點將沒用於輸出推理的Op節點從圖中剝離掉,再重新保存到指定的文件里(用write_graphdef或Saver)

文件目錄:tensorflow/python/tools/free_graph.py
測試文件:tensorflow/python/tools/free_graph_test.py 這個測試文件很有學習價值
參數:
總共有11個參數,一個個介紹下(必選: 表示必須有值;可選: 表示可以為空):
1、input_graph:(必選)模型文件,可以是二進制的pb文件,或文本的meta文件,用input_binary來指定區分(見下面說明)
2、input_saver:(可選)Saver解析器。保存模型和權限時,Saver也可以自身序列化保存,以便在加載時應用合適的版本。主要用於版本不兼容時使用。可以為空,為空時用當前版本的Saver。
3、input_binary:(可選)配合input_graph用,為true時,input_graph為二進制,為false時,input_graph為文件。默認False
4、input_checkpoint:(必選)檢查點數據文件。訓練時,給Saver用於保存權重、偏置等變量值。這時用於模型恢復變量值。
5、output_node_names:(必選)輸出節點的名字,有多個時用逗號分開。用於指定輸出節點,將沒有在輸出線上的其它節點剔除。
6、restore_op_name:(可選)從模型恢復節點的名字。升級版中已棄用。默認:save/restore_all
7、filename_tensor_name:(可選)已棄用。默認:save/Const:0
8、output_graph:(必選)用來保存整合后的模型輸出文件。
9、clear_devices:(可選),默認True。指定是否清除訓練時節點指定的運算設備(如cpu、gpu、tpu。cpu是默認)
10、initializer_nodes:(可選)默認空。權限加載后,可通過此參數來指定需要初始化的節點,用逗號分隔多個節點名字。
11、variable_names_blacklist:(可先)默認空。變量黑名單,用於指定不用恢復值的變量,用逗號分隔多個變量名字。
用法:
例: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,導入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文件

二、代碼:

https://blog.csdn.net/u012426298/article/details/85935946

兩步:

CKPT 轉換成 PB格式

pb模型預測

1、CKPT 轉換成 PB格式

 將CKPT 轉換成 PB格式的文件的過程可簡述如下:

通過傳入 CKPT 模型的路徑得到模型的圖和變量數據
通過 import_meta_graph 導入模型中的圖
通過 saver.restore 從模型中恢復圖中各個變量的數據
通過 graph_util.convert_variables_to_constants 將模型持久化【變成常量,這樣才能將變量和模型圖一起存儲】
 下面的CKPT 轉換成 PB格式例子,是我訓練GoogleNet InceptionV3模型保存的ckpt轉pb文件的例子,訓練過程可參考博客:《使用自己的數據集訓練GoogLenet InceptionNet V1 V2 V3模型(TensorFlow)》:

    def freeze_graph(input_checkpoint,output_graph):
        '''
        :param input_checkpoint:
        :param output_graph: PB模型保存路徑
        :return:
        '''
        # checkpoint = tf.train.get_checkpoint_state(model_folder) #檢查目錄下ckpt文件狀態是否可用
        # input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑
     
        # 指定輸出的節點名稱,該節點名稱必須是原模型中存在的節點
        output_node_names = "InceptionV3/Logits/SpatialSqueeze"
        saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
        graph = tf.get_default_graph() # 獲得默認的圖
        input_graph_def = graph.as_graph_def()  # 返回一個序列化的圖代表當前的圖
     
        with tf.Session() as sess:
            saver.restore(sess, input_checkpoint) #恢復圖並得到數據
            output_graph_def = graph_util.convert_variables_to_constants(  # 模型持久化,將變量值固定
                sess=sess,
                input_graph_def=input_graph_def,# 等於:sess.graph_def
                output_node_names=output_node_names.split(","))# 如果有多個輸出節點,以逗號隔開
     
            with tf.gfile.GFile(output_graph, "wb") as f: #保存模型
                f.write(output_graph_def.SerializeToString()) #序列化輸出
            print("%d ops in the final graph." % len(output_graph_def.node)) #得到當前圖有幾個操作節點
     
            # for op in graph.get_operations():
            #     print(op.name, op.values())

說明:

1、函數freeze_graph中,最重要的就是要確定“指定輸出的節點名稱”,這個節點名稱必須是原模型中存在的節點,對於freeze操作,我們需要定義輸出結點的名字。因為網絡其實是比較復雜的,定義了輸出結點的名字,那么freeze的時候就只把輸出該結點所需要的子圖都固化下來,其他無關的就舍棄掉。因為我們freeze模型的目的是接下來做預測。所以,output_node_names一般是網絡模型最后一層輸出的節點名稱,或者說就是我們預測的目標。

 2、在保存的時候,通過convert_variables_to_constants函數來指定需要固化的節點名稱,對於鄙人的代碼,需要固化的節點只有一個:output_node_names。注意節點名稱與張量的名稱的區別,例如:“input:0”是張量的名稱,而"input"表示的是節點的名稱。

3、源碼中通過graph = tf.get_default_graph()獲得默認的圖,這個圖就是由saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)恢復的圖,因此必須先執行tf.train.import_meta_graph,再執行tf.get_default_graph() 。

4、實質上,我們可以直接在恢復的會話sess中,獲得默認的網絡圖,更簡單的方法,如下:

def freeze_graph(input_checkpoint,output_graph):
    '''
    :param input_checkpoint:
    :param output_graph: PB模型保存路徑
    :return:
    '''
    # checkpoint = tf.train.get_checkpoint_state(model_folder) #檢查目錄下ckpt文件狀態是否可用
    # input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑
 
    # 指定輸出的節點名稱,該節點名稱必須是原模型中存在的節點
    output_node_names = "InceptionV3/Logits/SpatialSqueeze"
    saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
 
    with tf.Session() as sess:
        saver.restore(sess, input_checkpoint) #恢復圖並得到數據
        output_graph_def = graph_util.convert_variables_to_constants(  # 模型持久化,將變量值固定
            sess=sess,
            input_graph_def=sess.graph_def,# 等於:sess.graph_def
            output_node_names=output_node_names.split(","))# 如果有多個輸出節點,以逗號隔開
 
        with tf.gfile.GFile(output_graph, "wb") as f: #保存模型
            f.write(output_graph_def.SerializeToString()) #序列化輸出
        print("%d ops in the final graph." % len(output_graph_def.node)) #得到當前圖有幾個操作節點

調用freeze_graph方法:

調用方法很簡單,輸入ckpt模型路徑,輸出pb模型的路徑即可:

    # 輸入ckpt模型路徑
    input_checkpoint='models/model.ckpt-10000'
    # 輸出pb模型的路徑
    out_pb_path="models/pb/frozen_model.pb"
    # 調用freeze_graph將ckpt轉為pb
    freeze_graph(input_checkpoint,out_pb_path)
5、上面以及說明:在保存的時候,通過convert_variables_to_constants函數來指定需要固化的節點名稱,對於鄙人的代碼,需要固化的節點只有一個:output_node_names。因此,其他網絡模型,也可以通過簡單的修改輸出的節點名稱output_node_names,將ckpt轉為pb文件 。

       PS:注意節點名稱,應包含name_scope 和 variable_scope命名空間,並用“/”隔開,如"InceptionV3/Logits/SpatialSqueeze"

2、 pb模型預測
    下面是預測pb模型的代碼:

def freeze_graph_test(pb_path, image_path):
    '''
    :param pb_path:pb文件的路徑
    :param image_path:測試圖片的路徑
    :return:
    '''
    with tf.Graph().as_default():
        output_graph_def = tf.GraphDef()
        with open(pb_path, "rb") as f:
            output_graph_def.ParseFromString(f.read())
            tf.import_graph_def(output_graph_def, name="")
        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())
 
            # 定義輸入的張量名稱,對應網絡結構的輸入張量
            # input:0作為輸入圖像,keep_prob:0作為dropout的參數,測試時值為1,is_training:0訓練參數
            input_image_tensor = sess.graph.get_tensor_by_name("input:0")
            input_keep_prob_tensor = sess.graph.get_tensor_by_name("keep_prob:0")
            input_is_training_tensor = sess.graph.get_tensor_by_name("is_training:0")
 
            # 定義輸出的張量名稱
            output_tensor_name = sess.graph.get_tensor_by_name("InceptionV3/Logits/SpatialSqueeze:0")
 
            # 讀取測試圖片
            im=read_image(image_path,resize_height,resize_width,normalization=True)
            im=im[np.newaxis,:]
            # 測試讀出來的模型是否正確,注意這里傳入的是輸出和輸入節點的tensor的名字,不是操作節點的名字
            # out=sess.run("InceptionV3/Logits/SpatialSqueeze:0", feed_dict={'input:0': im,'keep_prob:0':1.0,'is_training:0':False})
            out=sess.run(output_tensor_name, feed_dict={input_image_tensor: im,
                                                        input_keep_prob_tensor:1.0,
                                                        input_is_training_tensor:False})
            print("out:{}".format(out))
            score = tf.nn.softmax(out, name='pre')
            class_id = tf.argmax(score, 1)
            print "pre class_id:{}".format(sess.run(class_id))

說明:

1、與ckpt預測不同的是,pb文件已經固化了網絡模型結構,因此,即使不知道原訓練模型(train)的源碼,我們也可以恢復網絡圖,並進行預測。恢復模型十分簡單,只需要從讀取的序列化數據中導入網絡結構即可:

tf.import_graph_def(output_graph_def, name="")
2、但必須知道原網絡模型的輸入和輸出的節點名稱(當然了,傳遞數據時,是通過輸入輸出的張量來完成的)。由於InceptionV3模型的輸入有三個節點,因此這里需要定義輸入的張量名稱,它對應網絡結構的輸入張量:

input_image_tensor = sess.graph.get_tensor_by_name("input:0")
input_keep_prob_tensor = sess.graph.get_tensor_by_name("keep_prob:0")
input_is_training_tensor = sess.graph.get_tensor_by_name("is_training:0")
以及輸出的張量名稱:

output_tensor_name = sess.graph.get_tensor_by_name("InceptionV3/Logits/SpatialSqueeze:0")
3、預測時,需要feed輸入數據:

# 測試讀出來的模型是否正確,注意這里傳入的是輸出和輸入節點的tensor的名字,不是操作節點的名字
# out=sess.run("InceptionV3/Logits/SpatialSqueeze:0", feed_dict={'input:0': im,'keep_prob:0':1.0,'is_training:0':False})
out=sess.run(output_tensor_name, feed_dict={input_image_tensor: im,
                                            input_keep_prob_tensor:1.0,
                                            input_is_training_tensor:False})
 4、其他網絡模型預測時,也可以通過修改輸入和輸出的張量的名稱 。

       PS:注意張量的名稱,即為:節點名稱+“:”+“id號”,如"InceptionV3/Logits/SpatialSqueeze:0"

 

完整的CKPT 轉換成 PB格式和預測的代碼如下:

# -*-coding: utf-8 -*-
"""
    @Project: tensorflow_models_nets
    @File   : convert_pb.py
    @Author : panjq
    @E-mail : pan_jinquan@163.com
    @Date   : 2018-08-29 17:46:50
    @info   :
    -通過傳入 CKPT 模型的路徑得到模型的圖和變量數據
    -通過 import_meta_graph 導入模型中的圖
    -通過 saver.restore 從模型中恢復圖中各個變量的數據
    -通過 graph_util.convert_variables_to_constants 將模型持久化
"""
 
import tensorflow as tf
from create_tf_record import *
from tensorflow.python.framework import graph_util
 
resize_height = 299  # 指定圖片高度
resize_width = 299  # 指定圖片寬度
depths = 3
 
def freeze_graph_test(pb_path, image_path):
    '''
    :param pb_path:pb文件的路徑
    :param image_path:測試圖片的路徑
    :return:
    '''
    with tf.Graph().as_default():
        output_graph_def = tf.GraphDef()
        with open(pb_path, "rb") as f:
            output_graph_def.ParseFromString(f.read())
            tf.import_graph_def(output_graph_def, name="")
        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())
 
            # 定義輸入的張量名稱,對應網絡結構的輸入張量
            # input:0作為輸入圖像,keep_prob:0作為dropout的參數,測試時值為1,is_training:0訓練參數
            input_image_tensor = sess.graph.get_tensor_by_name("input:0")
            input_keep_prob_tensor = sess.graph.get_tensor_by_name("keep_prob:0")
            input_is_training_tensor = sess.graph.get_tensor_by_name("is_training:0")
 
            # 定義輸出的張量名稱
            output_tensor_name = sess.graph.get_tensor_by_name("InceptionV3/Logits/SpatialSqueeze:0")
 
            # 讀取測試圖片
            im=read_image(image_path,resize_height,resize_width,normalization=True)
            im=im[np.newaxis,:]
            # 測試讀出來的模型是否正確,注意這里傳入的是輸出和輸入節點的tensor的名字,不是操作節點的名字
            # out=sess.run("InceptionV3/Logits/SpatialSqueeze:0", feed_dict={'input:0': im,'keep_prob:0':1.0,'is_training:0':False})
            out=sess.run(output_tensor_name, feed_dict={input_image_tensor: im,
                                                        input_keep_prob_tensor:1.0,
                                                        input_is_training_tensor:False})
            print("out:{}".format(out))
            score = tf.nn.softmax(out, name='pre')
            class_id = tf.argmax(score, 1)
            print "pre class_id:{}".format(sess.run(class_id))
 
 
def freeze_graph(input_checkpoint,output_graph):
    '''
    :param input_checkpoint:
    :param output_graph: PB模型保存路徑
    :return:
    '''
    # checkpoint = tf.train.get_checkpoint_state(model_folder) #檢查目錄下ckpt文件狀態是否可用
    # input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑
 
    # 指定輸出的節點名稱,該節點名稱必須是原模型中存在的節點
    output_node_names = "InceptionV3/Logits/SpatialSqueeze"
    saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
 
    with tf.Session() as sess:
        saver.restore(sess, input_checkpoint) #恢復圖並得到數據
        output_graph_def = graph_util.convert_variables_to_constants(  # 模型持久化,將變量值固定
            sess=sess,
            input_graph_def=sess.graph_def,# 等於:sess.graph_def
            output_node_names=output_node_names.split(","))# 如果有多個輸出節點,以逗號隔開
 
        with tf.gfile.GFile(output_graph, "wb") as f: #保存模型
            f.write(output_graph_def.SerializeToString()) #序列化輸出
        print("%d ops in the final graph." % len(output_graph_def.node)) #得到當前圖有幾個操作節點
 
        # for op in sess.graph.get_operations():
        #     print(op.name, op.values())
 
def freeze_graph2(input_checkpoint,output_graph):
    '''
    :param input_checkpoint:
    :param output_graph: PB模型保存路徑
    :return:
    '''
    # checkpoint = tf.train.get_checkpoint_state(model_folder) #檢查目錄下ckpt文件狀態是否可用
    # input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑
 
    # 指定輸出的節點名稱,該節點名稱必須是原模型中存在的節點
    output_node_names = "InceptionV3/Logits/SpatialSqueeze"
    saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
    graph = tf.get_default_graph() # 獲得默認的圖
    input_graph_def = graph.as_graph_def()  # 返回一個序列化的圖代表當前的圖
 
    with tf.Session() as sess:
        saver.restore(sess, input_checkpoint) #恢復圖並得到數據
        output_graph_def = graph_util.convert_variables_to_constants(  # 模型持久化,將變量值固定
            sess=sess,
            input_graph_def=input_graph_def,# 等於:sess.graph_def
            output_node_names=output_node_names.split(","))# 如果有多個輸出節點,以逗號隔開
 
        with tf.gfile.GFile(output_graph, "wb") as f: #保存模型
            f.write(output_graph_def.SerializeToString()) #序列化輸出
        print("%d ops in the final graph." % len(output_graph_def.node)) #得到當前圖有幾個操作節點
 
        # for op in graph.get_operations():
        #     print(op.name, op.values())
 
 
if __name__ == '__main__':
    # 輸入ckpt模型路徑
    input_checkpoint='models/model.ckpt-10000'
    # 輸出pb模型的路徑
    out_pb_path="models/pb/frozen_model.pb"
    # 調用freeze_graph將ckpt轉為pb
    freeze_graph(input_checkpoint,out_pb_path)
 
    # 測試pb模型
    image_path = 'test_image/animal.jpg'
    freeze_graph_test(pb_path=out_pb_path, image_path=image_path)

 


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM