tensorflow 之模型的保存与加载(一)

提示：扫一扫查出行
【扫一扫了解最新限行尾号】
复制提示

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

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

#!/usr/bin/env python3
#-*- coding:utf-8 -*-
############################
#File Name: saver.py
#Brief:
#Author: frank
#Mail: frank0903@aliyun.com
#Created Time:2018-06-22 22:12:52
############################ """
checkpoint #保存所有的模型文件列表
my_test_model.ckpt.data-00000-of-00001
my_test_model.ckpt.index
my_test_model.ckpt.meta #保存计算图的结构信息,即神经网络的结构
""" 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")
result = v1 + v2 init_op = tf.global_variables_initializer() #声明tf.train.Saver类用于保存模型
saver = tf.train.Saver() with tf.Session() as sess:
sess.run(init_op)
#将模型保存到指定路径
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")))

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

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