tensorflow模型文件pb、ckpt的生成、转换和加载

生成ckpt文件

# generate ckpt file
import tensorflow as tf
ckpt_file = "/home/xavier/temp/model"
x1 = tf.placeholder("float", name="x1")
x2 = tf.placeholder("float", name="x2")
w1 = tf.Variable(tf.random_normal(shape=[1], stddev=1.0, seed=1.0), name="w1")
w2 = tf.Variable(tf.random_normal(shape=[1], stddev=1.0, seed=1.0), name="w2")
sess = tf.Session()
result1 = tf.multiply(w1, x1, name="op_mul1")
result2 = tf.multiply(w2, x2, name="op_mul2")
result = tf.add(result1, result2, name="op_add1")
sess.run(tf.global_variables_initializer())
print(sess.run(result, feed_dict={x1:1.0, x2:10.0}))
saver = saver = tf.train.Saver()
saver.save(sess, ckpt_file)
# print result : [-8.9245]

加载ckpt文件

# load ckpt file
import tensorflow as tf
ckpt_meta_file = "/home/xavier/temp/model.meta"
ckpt_model_file = "/home/xavier/temp/model"
sess = tf.Session()
saver = tf.train.import_meta_graph(ckpt_meta_file)
saver.restore(sess, ckpt_model_file)
graph = tf.get_default_graph()
x1 = graph.get_tensor_by_name("x1:0")
x2 = graph.get_tensor_by_name("x2:0")
feed_dict = {x1:1.0, x2:10.0}
op_add = graph.get_tensor_by_name("op_add1:0")
print(sess.run(op_add, feed_dict))
# print result : [-8.9245]

ckpt转换成pb文件

# transform ckpt to pb file
import tensorflow as tf
from tensorflow.python.framework import graph_util
ckpt_meta_file = "/home/xavier/temp/model.meta"
ckpt_model_file = "/home/xavier/temp/model"
pb_file = "/home/xavier/temp/model.pb"
output_node_names = "op_add1"
saver = tf.train.import_meta_graph(ckpt_meta_file)
graph = tf.get_default_graph()
input_graph_def = graph.as_graph_def()
with tf.Session() as sess:
     saver.restore(sess, ckpt_model_file)
     output_graph_def = graph_util.convert_variables_to_constants(sess=sess, input_graph_def=input_graph_def, output_node_names=output_node_names.split(","))
     with tf.gfile.GFile(pb_file, "wb") as f:
         f.write(output_graph_def.SerializeToString())

加载pb文件（由＂ckpt转换成pb文件＂步骤生成的）

# load pb file
import tensorflow as tf
pb_file = "/home/xavier/temp/model.pb"
sess = tf.Session()
with tf.gfile.GFile(pb_file, "rb") as f:
     graph_def = tf.GraphDef()
     graph_def.ParseFromString(f.read())
     sess.graph.as_default()
     tf.import_graph_def(graph_def, name='')
sess.run(tf.global_variables_initializer())
print(sess.run("w1:0"))
print(sess.run("w2:0"))
x1 = sess.graph.get_tensor_by_name("x1:0")
x2 = sess.graph.get_tensor_by_name("x2:0")
op_add = sess.graph.get_tensor_by_name("op_add1:0")
print(sess.run(op_add, feed_dict={x1:1.0, x2:10.0}))
#print result : [-0.8113182] [-0.8113182] [-8.9245]

----------------------------------推荐方式--------------------------------------------

直接利用tf.saved_model生成pb文件

# generate pb file from tf.saved_model
import tensorflow as tf
pb_file = "/home/xavier/temp/saved_model"
ckpt_file = "/home/xavier/temp/model"
x1 = tf.placeholder("float", name="x1")
x2 = tf.placeholder("float", name="x2")
w1 = tf.Variable(tf.random_normal(shape=[1], stddev=1.0, seed=1.0), name="w1")
w2 = tf.Variable(tf.random_normal(shape=[1], stddev=1.0, seed=1.0), name="w2")
sess = tf.Session()
sess.graph.as_default()
result1 = tf.multiply(w1, x1, name="op_mul1")
result2 = tf.multiply(w2, x2, name="op_mul2")
result = tf.add(result1, result2, name="op_add1")
sess.run(tf.global_variables_initializer())
print(sess.run(result, feed_dict={x1:1.0, x2:10.0}))
builder = tf.saved_model.builder.SavedModelBuilder(pb_file)
builder.add_meta_graph_and_variables(sess, ["model1"])
builder.save()
# print result : [-8.9245]

加载pb文件（由上一步tf.saved_model模块生成）

# load pb file generated by tf.saved_model
import tensorflow as tf
pb_file = "/home/xavier/temp/saved_model"
sess = tf.Session()
tf.saved_model.loader.load(sess, ["model1"], pb_file)
sess.run(tf.global_variables_initializer())
x1 = sess.graph.get_tensor_by_name("x1:0")
x2 = sess.graph.get_tensor_by_name("x2:0")
op_add = sess.graph.get_tensor_by_name("op_add1:0")
print(sess.run("w1:0"))
print(sess.run("w2:0"))
print(sess.run(op_add, feed_dict={x1:1.0, x2:10.0}))
#print result : [-0.8113182] [-0.8113182] [-8.9245]

-------------------------------补充：如何通过模型获取张量名？----------------------------------------

上面的例子中都是通过记住张量名来加载模型并进行计算，那怎样才能直接通过模型获取张量名呢？可参考下面的例子：

# get tensor names of model
import tensorflow as tf
pb_file = "/home/xavier/temp/model.pb"
with tf.gfile.FastGFile(pb_file, "rb") as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    tf.import_graph_def(graph_def, name="")
tensor_name_list = [tensor.name for tensor in tf.get_default_graph().as_graph_def().node]
for tensor_name in tensor_name_list:
    print(tensor_name)

结果：

x1 
x2 
w1
w2 
op_mul1 
op_mul2 
op_add1 

参考文档：
https://blog.csdn.net/huachao1001/article/details/78501928
https://blog.csdn.net/guyuealian/article/details/82218092
https://blog.csdn.net/fu6543210/article/details/80343345
https://zhuanlan.zhihu.com/p/32887066
https://blog.csdn.net/ymznice/article/details/88826328