网络编程
位置:首页>> 网络编程>> Python编程>> TensorFlow实现保存训练模型为pd文件并恢复

TensorFlow实现保存训练模型为pd文件并恢复

作者:CrazyStoneZw  发布时间:2021-03-01 05:45:27 

标签:TensorFlow,训练模型,pd,恢复

TensorFlow保存模型代码


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/model.pb'
with tf.Session() as sess:
 sess.run(initop)
 print(sess.run(addop, feed_dict={x: 12, x2: 23}))
 output_graph_def = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['add', 'add2', 'add3'])
 # 将计算图写入到模型文件中
 model_f = tf.gfile.FastGFile(model_path, mode="wb")
 model_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}))

来源:https://blog.csdn.net/oYouHuo/article/details/81111833

0
投稿

猜你喜欢

手机版 网络编程 asp之家 www.aspxhome.com