【Tensorflow】模型保存与读取

Tensorflow模型结构

Meta graph:保存网络图结构，包括计算图中的变量，操作，集合，文件扩展名为.meta。

.data和.index文件：保存训练好的参数。

Checkpoint file:记录最新的模型。

保存模型

# 定义saver类
# 指定max_to_keep参数来选择保存最新的几个模型
# 设置keep_checkpoint_every_n_hours来指定隔几个小时保存一次模型
saver = tf.train.Saver(max_to_keep=4)
# 定义会话
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    ————训练步骤————
    # 保存模型，global_step参数会在命名时记录是在哪一步保存的模型
    # 在第一次保存模型时，图结构已经保存在.meta文件中，故之后不需要
    # 重复保存，可以设置write_meta_graph参数为False
    saver.save(sess, "model/my-model", global_step=epoch,global_step=step,write_meta_graph=False)

我们可以在tf.train.Saver()中指明想保存的参数列表，否则将自动保存模型所有参数。

import tensorflow as tf
w1 = tf.Variable(tf.random_normal(shape=[2]), name='w1')
w2 = tf.Variable(tf.random_normal(shape=[5]), name='w2')
saver = tf.train.Saver([w1,w2])
sess = tf.Session()
sess.run(tf.global_variables_initializer())
saver.save(sess, 'my_test_model',global_step=1000)

加载模型

若想使用一训练好的模型进行fine-tuning,需要两个步骤来加载模型：

1.构建网络图结构

可以通过手动写代码来构建网络图，也可以加载.meta文件来恢复图结构：

saver = tf.train.import_meta_graph('my_test_model-1000.meta')

2.载入训练好的参数

with tf.Session() as sess:
  new_saver = tf.train.import_meta_graph('my_test_model-1000.meta')
  new_saver.restore(sess, tf.train.latest_checkpoint('./'))