tensorflow 之模型的保存与加载

Posted 2020-11-20 black-mamba

怎样让通过训练的神经网络模型得以复用?

本文先介绍简单的模型保存与加载的方法，后续文章再慢慢深入解读.

 1 #!/usr/bin/env python3                                            
 2 #-*- coding:utf-8 -*-
 3 ############################
 4 #File Name: saver.py
 5 #Brief:
 6 #Author: frank
 7 #Mail: [email protected]
 8 #Created Time:2018-06-22 22:12:52
 9 ############################
10 
11 """
12 checkpoint                  #保存所有的模型文件列表
13 my_test_model.ckpt.data-00000-of-00001
14 my_test_model.ckpt.index
15 my_test_model.ckpt.meta     #保存计算图的结构信息,即神经网络的结构
16 """
17 
18 
19 import tensorflow as tf
20 
21 #声明两个变量并计算它们的和
22 v1 = tf.Variable(tf.constant(1.0, shape=[1]), name="v1")
23 v2 = tf.Variable(tf.constant(2.0, shape=[1]), name="v2")
24 result = v1 + v2
25 
26 init_op = tf.global_variables_initializer()
27 
28 #声明tf.train.Saver类用于保存模型
29 saver = tf.train.Saver()
30 
31 with tf.Session() as sess:
32     sess.run(init_op)
33     #将模型保存到指定路径
34     saver.save(sess,"my_test_model.ckpt")

 

模型的加载方法：

#!/usr/bin/env python3                                   
#-*- coding:utf-8 -*-                                    
############################                             
#File Name: restore.py                                   
#Brief:                                                  
#Author: frank                                           
#Mail: [email protected]                              
#Created Time:2018-06-22 22:34:16                        
############################                             
                                                         
                                                      
import tensorflow as tf                                  
                                                         
v1 = tf.Variable(tf.constant(1.0, shape=[1]), name="v1") 
v2 = tf.Variable(tf.constant(2.0, shape=[1]), name="v2") 
print(v1)                                                
result = v1 + v2                                         
print(result)                                            
                                                         
saver = tf.train.Saver()                                 
                                                         
with tf.Session() as sess:                               
    saver.restore(sess, "my_test_model.ckpt")            
    print(sess.run(result))                              
                                                         
                                                         
#运行结果:                                               
#<tf.Variable ‘v1:0‘ shape=(1,) dtype=float32_ref>       
#Tensor("add:0", shape=(1,), dtype=float32)              
#[3.]                                                    
                                                         
                                                      

上面的过程中还是定义了 图的结构，有点重复了，那么可不可以直接从已保存的ckpt中加载图呢?

import tensorflow as tf                                                
                                                                       
saver = tf.train.import_meta_graph("my_test_model.ckpt.meta")          
                                                                       
with tf.Session() as sess:                                             
    saver.restore(sess, "my_test_model.ckpt")                          
    print(sess.run(tf.get_default_graph().get_tensor_by_name("add:0")))

上面的程序，默认保存和加载了计算图中的全部变量,但有时可能只需要保存或加载部分变量。因为并不是所有隐藏层的参数需要重新训练。

具体怎么做呢?且听下回分解