TensorFlow拾遺(二) pb文件的生成

 

---

（一）獲取ckpt各節點名稱

方式一

方式二

（二）將ckpt轉化為pb文件

附錄：ckpt文件形式　　  

---

 　　將模型移植到諸如Android，FPGA等移動端時，需要模型的pb文件。深度學習框架會將模型權重保存為自身的格式，如.ckpt(tensorflow) .h5(tf.keras/Keras) .pt(pytorch)。此時，便需要對權重文件的格式進行轉換。

（一）獲取ckpt各節點名稱

　　在進行ckpt權重文件轉化為pb文件時，需要獲取到模型的輸出節點名稱，便於從相應的op中獲得op的輸出。

方式一

　　執行下列腳本，可以獲得模型各節點變量的txt文件。需要將checkpoint_path更改為自己的模型路徑，本篇使用的是SiamMask算法的模型。

# -*- coding: utf-8 -*-
# @Time    : 2020/7/14 下午3:11
# @Author  : monologuesmw   
# @Email   : [e-mail]
# @File    : node_name.py
# @Software: PyCharm
from tensorflow.python import pywrap_tensorflow
import os
def get_node_name(checkpoint_path):
    reader=pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
    var_to_shape_map=reader.get_variable_to_shape_map()
    node_name = []
    for key in var_to_shape_map:
        node_name.append(key)
    return node_name


def txt_save(filename, data):
    file = open(filename, 'a')
    for i in data:
        s = str(i) + '\n'
        file.write(s)
    file.close()


if __name__ == '__main__':
    checkpoint_path = 'ckpt/model.ckpt'
    file_name = 'node_name.txt'
    node_name = get_node_name(checkpoint_path)
    txt_save(file_name, node_name)

　　對這種方式獲得的節點名稱感興趣的可以猛戳這里。不過，從名稱上觀察，這種方式獲得的名稱是亂序的，不便於查找需要的輸出節點。

　　這種方式獲得的節點名稱如下所示

Score/Head/conv2d_1/bias
BBox/Head/conv2d_1/bias

方式二

　　還可以使用更簡潔的方式，獲取按順序排列的各種op：

import tensorflow as tf
def get_op(input_checkpoint):
    saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
    graph = tf.get_default_graph()
    txt_save_op(graph.get_operations(), "op_name")   # get node name

        
def txt_save_op(data, output_file):
    file = open(output_file, 'a')
    for op in data:
        s = str(op.name) + '\n'
        file.write(s)
    file.close()

if __name__ == '__main__':
    checkpoint_path = 'ckpt/model.ckpt'
    get_op(checkpoint_path)

　　這種方式會獲得tensorflow的所有op。由於op較多，感興趣的可以猛戳這里。

　　可以獲得兩個輸出節點的名稱：

BBox/Head/conv2d_1/BiasAdd
Score/Head/conv2d_1/BiasAdd

方式三　　

        另外，通過網絡結構變量的方式也可以獲取到網絡輸出節點的名稱，詳情猛戳這里。

（二）將ckpt轉化為pb文件

　　基於獲得的網絡輸出節點的名稱，便可以將ckpt轉換為pb。

　　核心代碼集中在freeze_graph()函數中。

　　如下述所示，值得注意的是代碼第15行，多個輸出節點需要用逗號隔開，並且逗號兩側不能有空格。

import tensorflow as tf
from tensorflow import graph_util
import os

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 = "BBox/Head/conv2d_1/BiasAdd,Score/Head/conv2d_1/BiasAdd"   # character , would't have space
    saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)

    # graph = tf.get_default_graph()
    # txt_save_op(graph.get_operations(), "op_name_1")   # get node name

    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())  # 序列化輸出


def txt_save(data, output_file):
    file = open(output_file, 'a')
    for i in data:
        s = str(i) + '\n'
        file.write(s)
    file.close()


def txt_save_op(data, output_file):
    file = open(output_file, 'a')
    for op in data:
        # s1 =str(op.name)
        # s2 = str(op.values)
        # s = s1+s2+'\n'
        s = str(op.name) + '\n'
        file.write(s)
    file.close()

# 獲取網絡結構變量  可省略
def network_param(input_checkpoint, output_file=None):
    saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
    with tf.Session() as sess:
        saver.restore(sess, input_checkpoint)
        variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
        txt_save(variables, output_file)


if __name__ == '__main__':
    checkpoint_path = 'ckpt/model.ckpt'
    output_path = 'pb/frozen_moedl.pb'
    a = os.path.split(output_path)[0]
    if not os.path.exists(a):
        os.makedirs(a)
    # 獲取網絡結構與變量  可省略
    #output_file = 'network_param.txt'
    #if not os.path.exists(output_file):
    #    network_param(checkpoint_path, output_file)
        
    freeze_graph(checkpoint_path, output_path)

　　其中，注釋掉的部分代碼是打印網絡結構與變量的，也可以獲得網絡輸出節點名稱，詳情猛戳這里。

附錄：ckpt文件形式

 

 

參考鏈接：https://blog.csdn.net/guyuealian/article/details/82218092#%C2%A0%20%C2%A0%20%C2%A01%E3%80%81%E8%AE%AD%E7%BB%83%E6%96%B9%E6%B3%95