tensorflow 之模型的保存与加载(一)
怎样让通过训练的神经网络模型得以复用?
本文先介绍简单的模型保存与加载的方法,后续文章再慢慢深入解读.
1 #!/usr/bin/env python3 2 #-*- coding:utf-8 -*- 3 ############################ 4 #File Name: saver.py 5 #Brief: 6 #Author: frank 7 #Mail: frank0903@aliyun.com 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: frank0903@aliyun.com #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")))
上面的程序,默认保存和加载了计算图中的全部变量,但有时可能只需要保存或加载部分变量。因为并不是所有隐藏层的参数需要重新训练。
具体怎么做呢?且听下回分解