tensorflow中保存部分模型节点

最新推荐文章于 2024-03-20 10:49:17 发布

原创最新推荐文章于 2024-03-20 10:49:17 发布 · 2.2k 阅读

0 ·

CC 4.0 BY-SA版权

tensorflow用法专栏收录该内容

64 篇文章

订阅专栏

本文介绍如何在TensorFlow中仅保存神经网络模型的部分节点，例如排除变量初始化等不必要的节点。示例中展示了保存包含特定节点的模型状态。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

比如在神经网络做预测的时候，只需要保存部分的节点，比如变量的初始化这些就不需要保存。

比如：

import tensorflow as tf;  
import numpy as np;  
import matplotlib.pyplot as plt;  
from tensorflow.python.framework import graph_util

v1 = tf.Variable(tf.constant(1.0, shape=[1]), name='v1')
v2 = tf.Variable(tf.constant(2.0, shape=[1]), name='v2')
result = v1 + v2

init = tf.initialize_all_variables()

with tf.Session() as sess:
	sess.run(init)

	graph_def = tf.get_default_graph().as_graph_def()

	output_graph_def = graph_util.convert_variables_to_constants(sess, graph_def, ['add'])
	with tf.gfile.GFile('/home/penglu/Desktop/lplp/model.pb', "wb") as f:
		f.write(output_graph_def.SerializeToString())

导入模型：

import tensorflow as tf;  
import numpy as np;  
import matplotlib.pyplot as plt;  
from tensorflow.python.framework import graph_util

with tf.Session() as sess:
	model_filename = '/home/penglu/Desktop/lplp/model.pb'

	with tf.gfile.FastGFile(model_filename, 'rb') as f:
		graph_def = tf.GraphDef()
		graph_def.ParseFromString(f.read())

		result = tf.import_graph_def(graph_def, return_elements=['add:0'])
		print sess.run(result)

输出：

[array([ 3.], dtype=float32)]