TensorFlow网络框架模型的保存和加载

在训练深度学习模型时， 需要对将模型的参数保存，以便以后的需要。以前训练深度学习模型主要关注网络的设计，对模型的保存和加载不太注意。最近写使用TensorFlow写深度学习框架，觉得有必要详细了解一下。

最简单的保存和恢复模型的方法是使用 tf.train.Saver 对象。构造器给graph的所有变量，或是定义在列表里的变量，添加 save 和 restore ops。saver对象提供了方法来运行这些ops，定义检查点文件的读写路径。这个链接对TensorFlow的加载和保存模型讲解的比较清楚：https://cv-tricks.com/tensorflow-tutorial/save-restore-tensorflow-models-quick-complete-tutorial/

保存变量

用 tf.train.Saver() 创建一个 Saver 来管理模型中的所有变量。

# Create some variables.

v1 = tf.Variable(..., name="v1")

v2 = tf.Variable(..., name="v2")

...

# Add an op to initialize the variables.

init_op = tf.initialize_all_variables()

# Add ops to save and restore all the variables.

saver = tf.train.Saver()

# Later, launch the model, initialize the variables, do some work, save the

# variables to disk.

with tf.Session() as sess:

sess.run(init_op)

# Do some work with the model.

..

# Save the variables to disk.

save_path = saver.save(sess, "/tmp/model.ckpt")

print "Model saved in file: ", save_path

 

恢复变量

用同一个 Saver 对象来恢复变量。注意，当你从文件中恢复变量时，不需要事先对它们做初始化。

# Create some variables.

v1 = tf.Variable(..., name="v1")

v2 = tf.Variable(..., name="v2")

...

# Add ops to save and restore all the variables.

saver = tf.train.Saver()

# Later, launch the model, use the saver to restore variables from disk, and

# do some work with the model.

with tf.Session() as sess:

# Restore variables from disk.

saver.restore(sess, "/tmp/model.ckpt")

print "Model restored."

# Do some work with the model

...