Tensorflow 中的模型保存及权重迁移

模型的保存与加载

sess = tf.Session()
saver = tf.train.Saver()
saver.save(sess, './model/test') #保存模型
saver.restore(sess, 'name')

通过以上代码可以保存一个模型，一般生成 四个文件

第一个记录了最新保存的模型路径，第二个保存了训练好的权值，第三个是一些键值信息，第四个是模型的grah。
代码中restore是权值的重新加载，经常用于接力训练。

但是，如果我们训练好一个网络之后，仅仅想用这个网络的前几层的权值作为另一个网络的初始化的预训练权值该怎么办呢？

有一种方案：按照原来的网络结构，重新写一遍代码，然后将新的网络结构代码也写在后面，很显然这种方式能够能够完成目标，但是不够nice。
第二种方案：将训练的第一个网络的权重全部sess.run成真实的值存为 .mat。然后在新的网络中，将这.mat初始化进网络，或者assign进对应的权重位置，但这种方式很繁琐，工作量较大。
其实，上面我们提出的问题，就是预训练权重的迁移问题，这在迁移学习经常出现。在tensorflow中早有函数解决的了这个问题，代码如下：

1. 预训练网络

import tensorflow as tf

X = tf.placeholder(tf.float32, shape=[None, 2], name='X')
Y = tf.placeholder(tf.float32, shape=[None, 1], name='Y')
W = tf.get_variable("w", shape=[2, 1], initializer=tf.constant_initializer([0.1, 0.2]))
b = tf.get_variable("b", shape=[1], initializer=tf.constant_initializer(0.1))
hypothesis = tf.matmul(X, W) + b
cost = -tf.reduce_mean(Y * tf.log(hypothesis) + (1 - Y) * tf.log(1 - hypothesis))
optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.01)
train = optimizer.minimize(cost)
tf.add_to_collection('X', X) # 输入
tf.add_to_collection('Y', Y) # 标签
tf.add_to_collection('hypothesis', hypothesis) # 中间的运算结果

saver = tf.train.Saver()
sess = tf.Session()
sess.run(tf.global_variables_initializer())
saver.save(sess, './model/test')  #保存模型

第一个网络中，用原始的方法保存了网络模型，只是利用tf.add_to_collection()函数将网络的输入，以及我们需要利用的中间层hypothesis单独保存了起来。
2. 迁移训练的网络

import tensorflow as tf

sess = tf.Session()
grah_saver = tf.train.import_meta_graph('./model/test.meta')  # 加载模型的grah
grah_saver.restore(sess, './model/test')  # 加载模型中各种变量的值，注意这里不用文件的后缀
X = tf.get_collection("X")[0]
Y = tf.get_collection("Y")[0]
hypothesis = tf.get_collection("hypothesis")[0]
b = tf.get_variable("b1", shape=[1, 1], initializer=tf.constant_initializer(0.0))
new_layers = hypothesis + b
cost = -tf.reduce_mean(Y * tf.log(new_layers) + (1 - Y) * tf.log(1 - new_layers))
sess.run(tf.global_variables_initializer())
optm = tf.train.GradientDescentOptimizer(learning_rate=0.001).minimize(cost)
for _ in range(10):
    _, loss = sess.run([optm, cost], feed_dict={X: [[1, 2]], Y: [[1]]})
    print(loss)

在第2个模型中我们可以看见，通过grah_saver = tf.train.import_meta_graph(’./model/test.meta’) 直接加载了原始模型的grah，而grah_saver.restore(sess, ‘./model/test’) 加载了grah中的权值。对于原始网络的输入XY，以及中间层hypothesis，利用tf.get_collection()进行了获取。这样后面的新的网络层new_layers直接作用在hypothesis上，通过获取的XY进行输入，便可以进行权重的迁移训练。若还有疑问的同学可以自行调试这两个代码，感受一下即可。