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