tensorflow inception_resnet_v2 ckpt2pb, 自定义输入输出节点

获取ckpt模型中的节点名称

from tensorflow.python import pywrap_tensorflow
checkpoint_path = 'model.ckpt-xxx'
reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
var_to_shape_map = reader.get_variable_to_shape_map()
for key in var_to_shape_map:
    print("tensor_name: ", key)

将ckpt模型转为pb模型

import sys
import tensorflow as tf
from nets import nets_factory

input_node = tf.placeholder(tf.float32, [1, 299, 299, 3], name='input')
output_node_names = 'InceptionResnetV2/Logits/Predictions'
network_fn = nets_factory.get_network_fn('inception_resnet_v2', num_classes=4, is_training=False)
ckpt = sys.argv[1]
pb = sys.argv[2]

with tf.Session() as sess:
    endpoint = network_fn(input_node)
    sess.run(tf.global_variables_initializer())

    input_graph_def = tf.get_default_graph().as_graph_def()
    # node_names = [n.name for n  in input_graph_def.node]
    # for node in node_names:
    #     print (node)
    output_graph_def = tf.graph_util.convert_variables_to_constants(
        sess,  # The session
        input_graph_def,  # input_graph_def is useful for retrieving the nodes
        output_node_names.split(",")
    )

with tf.gfile.GFile(pb, 'wb') as f:
    f.write(output_graph_def.SerializeToString())