tensorflow对pd文件的保存和加载

代码只是一份笔记，以免以后忘记了

import tensorflow as tf

from tensorflow.python.framework import graph_util

var1 = tf.Variable(1.0, dtype=tf.float32, name='v1')

var2 = tf.Variable(2.0, dtype=tf.float32, name='v2')

var3 = tf.Variable(2.0, dtype=tf.float32, name='v3')

x = tf.placeholder(dtype=tf.float32, shape=None, name='x')

x2 = tf.placeholder(dtype=tf.float32, shape=None, name='x2')

addop = tf.add(x, x2, name='add')

addop2 = tf.add(var1, var2, name='add2')

addop3 = tf.add(var3, var2, name='add3')

initop = tf.global_variables_initializer()

model_path = './Test/model2.pb'

with tf.Session() as sess:

    sess.run(initop)

    print(sess.run(addop, feed_dict={x: 12, x2: 23}))  # 35

    output_graph_def = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['add', 'add2', 'add3'])

    # 将计算图写入到模型文件中
    with tf.gfile.FastGFile(model_path, mode="wb") as f:
        f.write(output_graph_def.SerializeToString())

重新加载模型图和模型参量

import tensorflow as tf

with tf.Session() as sess:

    model_f = tf.gfile.FastGFile("./Test/model.pb", mode='rb')

    graph_def = tf.GraphDef()

    graph_def.ParseFromString(model_f.read())

    c = tf.import_graph_def(graph_def, return_elements=["add2:0"])

    c2 = tf.import_graph_def(graph_def, return_elements=["add3:0"])

    x, x2, c3 = tf.import_graph_def(graph_def, return_elements=["x:0", "x2:0", "add:0"])


    print(sess.run(c))
    print(sess.run(c2))
    print(sess.run(c3, feed_dict={x: 23, x2: 2}))