tensorflow模型ckpt转pb

最后一个节点是 'save/restore_all'? 'Merge/MergeSummary' ?

import tensorflow as tf
from tensorflow.python.framework import graph_util

def ckpt2pb():
    with tf.Graph().as_default() as graph_old:
        isess = tf.InteractiveSession()

        ckpt_filename = 'model.ckpt-101'
        isess.run(tf.global_variables_initializer())
        saver = tf.train.import_meta_graph(ckpt_filename+'.meta', clear_devices=True)
        saver.restore(isess, ckpt_filename)
        op = tf.get_default_graph().get_operations()
        print(op[5650:-1])

        constant_graph = tf.compat.v1.graph_util.convert_variables_to_constants(isess, isess.graph_def, ['save/restore_all'])
        constant_graph = tf.compat.v1.graph_util.remove_training_nodes(constant_graph)

        with tf.io.gfile.GFile('model.pb', mode='wb') as f:
            f.write(constant_graph.SerializeToString())
        
ckpt2pb()

Reference

tensorflow模型ckpt转pb

发布了1636 篇原创文章 · 获赞 341 · 访问量 221万+

猜你喜欢

转载自blog.csdn.net/tony2278/article/details/103934779