tensorflow保存部分变量

tensorflow模型保存函数为:

saver = tf.train.Saver()

查看该函数初始化函数,输入参数为:

def __init__(self,
var_list=None,
reshape=False,
sharded=False,
max_to_keep=5,
keep_checkpoint_every_n_hours=10000.0,
name=None,
restore_sequentially=False,
saver_def=None,
builder=None,
defer_build=False,
allow_empty=False,
write_version=saver_pb2.SaverDef.V2,
pad_step_number=False,
save_relative_paths=False):

var_list参数为我们需要保存的变量数组,如果不输入var_list,则默认保存所有的变量.例如下面代码:

import tensorflow as tf

v1= tf.Variable(tf.random_normal([784, 200], stddev=0.35), name="v1")
v2= tf.Variable(tf.zeros([200]), name="v2")

saver = tf.train.Saver()

with tf.Session() as sess:
init_op = tf.global_variables_initializer()
sess.run(init_op)
saver.save(sess,"checkpoint/model_test",global_step=1)

当我们需要保存部分变量时,我们可以定义一个需要保存的变量数组,例如我们只想保存变量v2,代码如下:

import tensorflow as tf

v1= tf.Variable(tf.random_normal([784, 200], stddev=0.35), name="v1")
v2= tf.Variable(tf.zeros([200]), name="v2")

saver = tf.train.Saver( [v2])
# saver = tf.train.Saver()
with tf.Session() as sess:
init_op = tf.global_variables_initializer()
sess.run(init_op)
saver.save(sess,"checkpoint/model_test",global_step=1)

需要注意的是,模型保存需要在某一目录下,例如上面代码保存在checkpoint目录下,模型名为model_test,如果保存在当前目录下,如下代码:

import tensorflow as tf

v1= tf.Variable(tf.random_normal([784, 200], stddev=0.35), name="v1")
v2= tf.Variable(tf.zeros([200]), name="v2")

saver = tf.train.Saver( [v2])
# saver = tf.train.Saver()
with tf.Session() as sess:
init_op = tf.global_variables_initializer()
sess.run(init_op)
saver.save(sess,"model_test",global_step=1)

则会报错:

Traceback (most recent call last):

  File "/home/qinghua/pythonWork/pix2pix-tensorflow/demo_model_part.py", line 11, in <module>

    saver.save(sess,"model_test",global_step=1)

  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/training/saver.py", line 1488, in save

    raise exc

tensorflow.python.framework.errors_impl.FailedPreconditionError: checkpoint.tmp2059e44e2c59453c824e5784c5bbef8e

Process finished with exit code 1