TensorFlow模型持久化

比眉伴天荒 2022-10-13 12:32 147阅读 0赞

> --------------------
> 
> 模型持久化的目的在于可以使模型训练后的结果重复使用，节省重复训练模型的时间。

# 模型保存 #

train.Saver类是TensorFlow提供的用于保存和还原模型的API，使用非常简单。

import tensorflow as tf
    
    
    # 声明两个变量并计算其加和
    a = tf.Variable(tf.constant([1.0, 2.0], shape=[2]), name='a')
    b = tf.Variable(tf.constant([3.0, 4.0], shape=[2]), name='b')
    result = a + b
    
    # 初始化全部变量的操作
    init_op = tf.global_variables_initializer()
    # 定义 Saver 类对象用于保存模型
    saver = tf.train.Saver()
    
    with tf.Session() as sess:
        sess.run(init_op)
        saver.save(sess, "./model/model.ckpt")

上面的代码实现了一个简单的TensorFlow模型持久化的功能。

save()函数的sess参数用于指定要保存的模型会话，save\_path参数用于指定路径。

通过Saver类的save()函数将TensorFlow模型保存到一个指定路径下的model.ckpt文件中。

（TensorFlow模型一般会保存在文件名为.ckpt的文件中，可以省略后缀名，但是好的编程习惯是对其加以指定）

虽然上面的程序只制定了一个文件路径，但是在这个文件目录下回出现4个文件：  
![在这里插入图片描述][20210706204803779.png]

*  checkpoint文件是一个文本文件，保存了一个目录下所有的模型文件列表。该文件会被自动更新，当有更多模型被保存到model目录下时，文件内容会更新为最新的训练模型。
 *  model.ckpt.data-00000-of-00001文件是一个二进制文件，保存了TensorFlow中每一个变量的取值。
 *  model.ckpt.index文件是一个二进制文件，保存了每一个变量的名称，是一个string-string的table，其中table的key值为tensor名，value值为BundleEntryProto。
 *  model.ckpt.meta文件是一个二进制文件，保存了计算图的结构。

将一个模型文件分成多个文件保存的原因是TensorFlow会将模型的计算图结构以及参数的取值分开来保存。

# 模型加载 #

TensorFlow也提供了相应的函数来加载保存的模型。

with tf.Session() as sess:
        saver.restore(sess, "./model/model.ckpt")
        print(sess.run(result))

输出：  
![在这里插入图片描述][20210706205850285.png]  
加载模型的代码和保存模型的代码相似，但是省略了初始化全部变量的过程。

使用restore()函数需要在模型参数恢复前定义计算图上的所有运算，并且变量名需要与模型存在的变量名保持一致，这样就可以将变量的值通过已保存的模型加载进来。

有时我们可能不希望重复定义计算图上的计算，太繁琐了，TensorFlow提供了import\_meta\_graph()函数加载模型的计算图。

import\_meta\_graph()函数的输入参数为.meta文件的路径，返回一个Saver类实例，再调用这个实例的restore()函数就可以恢复参数了。

saver = tf.train.import_meta_graph("./model/model.ckpt.meta")
    
    with tf.Session() as sess:
        saver.restore(sess, "./model/model.ckpt")
        # 获取默认计算图上指定节点处的张量
        print(sess.run(tf.get_default_graph().get_tensor_by_name("add:0")))

输出：  
![在这里插入图片描述][20210706215626599.png]  
.ckpt.meta文件保存了计算图的结构，通过import\_meta\_graph()函数将计算图导入到程序中并传递给saver，之后在会话中通过restore()函数对该计算图中变量的值进行加载。

get\_tensor\_by\_name()函数用于获取指定节点处的张量(add:0 表示add节点的第一个输出)。

[20210706204803779.png]: /images/20221005/6b8f4f49dc3044e28669ce23fd044e7a.png
[20210706205850285.png]: /images/20221005/3329fac519f94618822430189c53d292.png
[20210706215626599.png]: /images/20221005/78081a200548403ab9dae6d71fd0eb09.png