Tensorflow 模型保存与调用

Tensorflow 两种保存模型的方式：pb 和  saved_model 都可以。

 

1、pb

1.1 模型保存成pb

freozen_pb.py

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



with tf.Session(graph=tf.Graph()) as sess:
    x = tf.placeholder(tf.int32, name='in_x')
    y = tf.placeholder(tf.int32, name='in_y')
    b = tf.Variable(1, name='b')
    m = tf.multiply(x, y)
    a = tf.add(m, b, name='out_add')

    sess.run(tf.global_variables_initializer())
    constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['out_add'])

    feed_dict = {x: 10, y: 3}
    print(sess.run(a, feed_dict))

    with tf.gfile.FastGFile('./model.pb', mode='wb') as f:
        f.write(constant_graph.SerializeToString())

 

 

1.2 调用pb模型

call_pb.py

import tensorflow as tf
from tensorflow.python.platform import gfile


sess = tf.Session()
with gfile.FastGFile('./model.pb', '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('b:0'))

in_x = sess.graph.get_tensor_by_name('in_x:0')
in_y = sess.graph.get_tensor_by_name('in_y:0')
out_add = sess.graph.get_tensor_by_name('out_add:0')

ret = sess.run(out_add, feed_dict={in_x: 8, in_y: 9})
print(ret)

 

2、 saved_model

2.1 模型保存成saved model

freozen_sm.py

import os
import tensorflow as tf

saved_model_path = os.getcwd()

with tf.Session(graph=tf.Graph()) as sess:
    x = tf.placeholder(tf.int32, name='in_x')
    y = tf.placeholder(tf.int32, name='in_y')
    b = tf.Variable(1, name='b')
    m = tf.multiply(x, y)
    a = tf.add(m, b, name='out_add')

    sess.run(tf.global_variables_initializer())

    tf.saved_model.simple_save(sess, './sm', {'in_x': x, 'in_y': y}, {'out_add': a}, )

 

2.2 调用saved model模型

call_sm.py

import tensorflow as tf

sess = tf.Session()
tf.saved_model.load(sess, [tf.saved_model.tag_constants.SERVING], './sm')
in_x = sess.graph.get_tensor_by_name('in_x:0')
in_y = sess.graph.get_tensor_by_name('in_y:0')
out_add = sess.graph.get_tensor_by_name('out_add:0')

ret = sess.run(out_add, feed_dict={in_x: 8, in_y: 5})
print(ret)