1. 程式人生 > 其它 >深度學習之格式轉換筆記(二):CKPT 轉換成 PB格式檔案

深度學習之格式轉換筆記(二):CKPT 轉換成 PB格式檔案

技術標籤:深度學習pythontensorflow深度學習

我們使用tf.train.saver()儲存模型時會產生多個檔案,也就是說把計算圖的結構和圖上引數取值分成了不同的檔案儲存。這也是在tensorflow中常用的儲存方式。

儲存檔案的程式碼:

import tensorflow as tf
# 宣告兩個變數
v1 = tf.Variable(tf.random_normal([1, 2]), name="v1")
v2 = tf.Variable(tf.random_normal([2, 3]), name="v2")
init_op =
tf.global_variables_initializer() # 初始化全部變數 saver = tf.train.Saver() # 宣告tf.train.Saver類用於儲存模型 with tf.Session() as sess: sess.run(init_op) print("v1:", sess.run(v1)) # 列印v1、v2的值一會讀取之後對比 print("v2:", sess.run(v2)) saver_path = saver.save(sess, "save/model.ckpt-510"
) # 將模型儲存到save/model.ckpt-510檔案 print("Model saved in file:", saver_path)

這時候我們就可以看到結果
在這裡插入圖片描述
其中

  • checkpoint:檢查點檔案,檔案儲存了一個目錄下所有的模型檔案列表;
  • model.ckpt-510.meta:儲存了TensorFlow計算圖的結構,可以理解為神經網路的網路結構,該檔案可以被
    tf.train.import_meta_graph 載入到當前預設的圖來使用。
  • ckpt-510.data : 儲存模型中每個變數的取值
  • ckpt-510.index:可能是內部需要的某種索引來正確對映前兩個檔案,它通常不是必需的

真正部署的時候,一般人家不會給你ckpt模型的,而是固化成pb模型以後再給你用,現在我們就來看看怎麼將ckpt固化成pb模型。

實際完整程式碼:

# -*-coding: utf-8 -*-
import os
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())

            # 定義輸入的張量名稱,對應網路結構的輸入張量,往往是通過tf.placeholder呼叫的。
            # 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 = 'D:/pycharm/CarPlateIdentity-master/carIdentityData/model1/char_recongnize/model.ckpt-510'
    # 輸出pb模型的路徑
    out_dirpath = 'D:/pycharm/CarPlateIdentity-master/carIdentityData/model1/char_recongnize/pb/'
    os.makedirs(os.path.dirname(out_dirpath),exist_ok=True)
    out_pb_path = out_dirpath+"frozen_model.pb"
    # 呼叫freeze_graph將ckpt轉為pb
    freeze_graph(input_checkpoint, out_pb_path)
    print("the success cover")
    # 測試pb模型
    # image_path = 'test_image/animal.jpg'
    # freeze_graph_test(pb_path=out_pb_path, image_path=image_path)

在將ckpt轉換為pd過程中,會依據輸出節點來丟棄那些與輸出節點無關的引數,只保留與輸出節點存在上下文關係的引數,這也就是生成pd檔案的意義所在,即通過減少引數量降低模型的大小,所以在生成pd的過程中需要明確指定輸出節點是誰,這樣才能確定其依賴的需要固化的上下文引數。