Tensorflow提供了一个非常简单的API来实现保存和还原一个神经网络模型。这个API就是tf.train.Saver()类。以下代码给出了保存计算图的方法。
保存模型:
import tensorflow as tf import tensorflow as tf v1 = tf.Variable(tf.constant(1.0,shape = [1])) v2 = tf.Variable(tf.constant(2.0,shape = [1])) result= v1 + v2 saver = tf.train.Saver() init_op = tf.initialize_all_variables() with tf.Session()as sess: sess.run(init_op) saver.save(sess,'model / model.ckpt') saver.save(sess,'model / model.ckpt')
加载模型:
import tensorflow as tf v1 = tf.Variable(tf.constant(1.0,shape = [1])) v2 = tf.Variable(tf.constant(2.0,shape = [1])) result = v1 + v2 saver = tf.train.Saver() with tf.Session()as sess: saver.restore(SESS,“model / model.ckpt) print(sess.run(result))
在加载模型的代码中没有运行变量初始化,而是将变量的值通过模型保存的值加载进来如果不想重复定义图上的运算,也可以直接加载已经持久化的图代码如下:
import tensorflow as tf saver = tf.train.import_meta_graph('/ model / model.ckpt') with tf.Session()as sess: saver.restore(sess,'model / model.ckpt')
在上面的代码中,默认加载了计算图中的所有变量。但有时候可能只需要加载部分变量。比如,可能有一个之前训练好的五层神经网络模型,现在想尝试一个六层的神经网络,那么可以将前面五层神经网络中的参数直接加载到新模型,而仅仅将最后一层神经网络重新训练。
为了保存或加载部分变量,在声明tf.train.Saver类时可以提供一个列表来指定需要加载或保存的变量。比如,在上面的加载模型的代码中使用保护程序= tf.train.Saver([ V1])命令来构建tf.train.Saver类,那么只有变量V1会被加载进来。如果运行了修改只加载V1的代码,会得到变量为初始化的错误,因为V2没有被加载,所以V2运行初始化之前是没有值的。