TensorFlow .ckpt 模型的保存与恢复

.ckpt

.ckpt 全称为 checkpoint，代表着一个检查点，即为 model 训练过程中的一个快照，可能是在训练开始，也可能是在训练完成。

tensorflow新版本不会生成.ckpt文件，你只要将这四个文件放入一个文件夹并命名，测试时直接调用这个文件夹就行了，这就相当于旧版本的.ckpt文件。

一、TensorFlow常规模型加载方法

保存模型

tf.train.Saver()类，.save(sess, ckpt文件目录)方法

var_list是字典形式{变量名字符串: 变量符号}，相对应的restore也根据同样形式的字典将ckpt中的字符串对应的变量加载给程序中的符号。

加载模型

当我们基于checkpoint文件(ckpt)加载参数时，实际上我们使用Saver.restore取代了initializer的初始化

checkpoint文件会记录保存信息，通过它可以定位最新保存的模型：

ckpt = tf.train.get_checkpoint_state('./model/')
print(ckpt.model_checkpoint_path)

.meta文件保存了当前图结构

.index文件保存了当前参数名

.data文件保存了当前参数值

1. 不加载图结构，只加载参数

由于实际上我们参数保存的都是Variable变量的值，所以其他的参数值（例如batch_size）等，我们在restore时可能希望修改，但是图结构在train时一般就已经确定了，所以我们可以使用tf.Graph().as_default()新建一个默认图（建议使用上下文环境），利用这个新图修改和变量无关的参值大小，从而达到目的。

'''
使用原网络保存的模型加载到自己重新定义的图上
可以使用python变量名加载模型，也可以使用节点名
'''
import AlexNet as Net
import AlexNet_train as train
import random
import tensorflow as tf
 
IMAGE_PATH = './flower_photos/daisy/5673728_71b8cb57eb.jpg'
 
with tf.Graph().as_default() as g:
 
    x = tf.placeholder(tf.float32, [1, train.INPUT_SIZE[0], train.INPUT_SIZE[1], 3])
    y = Net.inference_1(x, N_CLASS=5, train=False)
 
    with tf.Session() as sess:
        # 程序前面得有 Variable 供 save or restore 才不报错
        # 否则会提示没有可保存的变量
        saver = tf.train.Saver()
 
        ckpt = tf.train.get_checkpoint_state('./model/')
        img_raw = tf.gfile.FastGFile(IMAGE_PATH, 'rb').read()
        img = sess.run(tf.expand_dims(tf.image.resize_images(
            tf.image.decode_jpeg(img_raw),[224,224],method=random.randint(0,3)),0))
 
        if ckpt and ckpt.model_checkpoint_path:
            print(ckpt.model_checkpoint_path)
            saver.restore(sess,'./model/model.ckpt-0')
            global_step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
            res = sess.run(y, feed_dict={x: img})
            print(global_step,sess.run(tf.argmax(res,1)))

2. 加载图结构和参数

'''
直接使用使用保存好的图
无需加载python定义的结构，直接使用节点名称加载模型
由于节点形状已经定下来了，所以有不便之处，placeholder定义batch后单张传会报错
现阶段不推荐使用，以后如果理解深入了可能会找到使用方法
'''
import AlexNet_train as train
import random
import tensorflow as tf
 
IMAGE_PATH = './flower_photos/daisy/5673728_71b8cb57eb.jpg'
 
 
ckpt = tf.train.get_checkpoint_state('./model/')                          # 通过检查点文件锁定最新的模型
saver = tf.train.import_meta_graph(ckpt.model_checkpoint_path +'.meta')   # 载入图结构，保存在.meta文件中
 
with tf.Session() as sess:
    saver.restore(sess,ckpt.model_checkpoint_path)                        # 载入参数，参数保存在两个文件中，不过restore会自己寻找
 
    img_raw = tf.gfile.FastGFile(IMAGE_PATH, 'rb').read()
    img = sess.run(tf.image.resize_images(
        tf.image.decode_jpeg(img_raw), train.INPUT_SIZE, method=random.randint(0, 3)))
    imgs = []
    for i in range(128):
       imgs.append(img)
    print(sess.run(tf.get_default_graph().get_tensor_by_name('fc3:0'),feed_dict={'Placeholder:0': imgs}))
 
    '''
    img = sess.run(tf.expand_dims(tf.image.resize_images(
        tf.image.decode_jpeg(img_raw), train.INPUT_SIZE, method=random.randint(0, 3)), 0))
    print(img)
    imgs = []
    for i in range(128):
        imgs.append(img)
    print(sess.run(tf.get_default_graph().get_tensor_by_name('conv1:0'),
                   feed_dict={'Placeholder:0':img}))

注意，在所有两种方式中都可以通过调用节点名称使用节点输出张量，节点.name属性返回节点名称。

3. 简化版本

# 连同图结构一同加载
ckpt = tf.train.get_checkpoint_state('./model/')
saver = tf.train.import_meta_graph(ckpt.model_checkpoint_path +'.meta')
with tf.Session() as sess:
    saver.restore(sess,ckpt.model_checkpoint_path)
             
# 只加载数据，不加载图结构，可以在新图中改变batch_size等的值
# 不过需要注意，Saver对象实例化之前需要定义好新的图结构，否则会报错
saver = tf.train.Saver()
with tf.Session() as sess:
    ckpt = tf.train.get_checkpoint_state('./model/')
    saver.restore(sess,ckpt.model_checkpoint_path)

参考：叠加态的猫-『TensorFlow』模型载入方法汇总