本帖最后由 keer_zu 于 2018-12-28 14:24 编辑
TensorFlow保存和恢复模型的方法总结
@sherwin @gaoyang9992006 @紫剑 @yyy71cj @happy_ice
@21ic小喇叭 @21ic小管家
使用TensorFlow训练模型的过程中，需要适时对模型进行保存，以及对保存的模型进行restore，以方便后续对模型进行处理。比如进行测试，或者部署；比如拿别的模型进行fine-tune，等等。当然，直接的保存和restore比较简单，无需多言，但是保存和restore中还牵涉到其他问题，以及针对各种需求的各种参数等，可能不便一下都记好。因此，有必要对此进行一个总结。本文就是对使用TensorFlow保存和restore模型的相关内容进行一下总结，以便备忘。
保存模型保存模型是整个内容的第一步，当然也十分简单。无非是创建一个saver，并在一个Session里完成保存。比如：
saver = tf.train.Saver()
with tf.Session() as sess:
saver.save(sess, model_name)
以上代码在0.11以下版本的TensorFlow里会保存与下面类似的3个文件：
checkpoint
model.ckpt-1000.meta
model.ckpt-1000.ckpt
在0.11及以上版本的TensorFlow里则会保存与下类似的4个文件：
checkpoint
model.ckpt-1000.index
model.ckpt-1000.data-00000-of-00001
model.ckpt-1000.meta
其中checkpoint列出保存的所有模型以及最近的模型；meta文件是模型定义的内容；ckpt（或data和index）文件是保存的模型数据；内里细节无需过多关注，如果想了解，stackOverflow上有一个解释的回答。
当然，除了上面最简单的保存方式，也可以指定保存的步数，多长时间保存一次，磁盘上最多保有几个模型（将前面的删除以保持固定个数），如下：
创建saver时指定参数：
saver = tf.train.Saver(savable_variables, max_to_keep=n, keep_checkpoint_every_n_hours=m)
其中savable_variables指定待保存的变量，比如指定为tf.global_variables()保存所有global变量；指定为[v1, v2]保存v1和v2两个变量；如果省略，则保存所有；
max_to_keep指定磁盘上最多保有几个模型；keep_checkpoint_every_n_hours指定多少小时保存一次。
保存模型时指定参数：
saver.save(sess, 'model_name', global_step=step,write_meta_graph=False)
如上，其中可以指定模型文件名，步数，write_meta_graph则用来指定是否保存meta文件记录graph等等。
Restore模型具体来说，Restore模型的过程可以分为两个部分，首先是创建模型，可以手动创建，也可以从meta文件里加载graph进行创建。
创建模型与训练模型时创建模型的代码相同，可以直接复制过来使用。
从meta文件里进行加载，可以直接在Session里进行如下操作：
with tf.Session() as sess:
saver = tf.train.import_meta_graph('model.ckpt-1000.meta')
后面的参数直接使用meta文件的路径即可。如此，即将模型定义的graph加载进来了。
当然，还有一点需要注意，并非所有的TensorFlow模型都能将graph输出到meta文件中或者从meta文件中加载进来，如果模型有部分不能序列化的部分，则此种方法可能会无效。
然后就是为模型加载数据，可以使用下面两种方法：
with tf.Session() as sess:
saver = tf.train.import_meta_graph('model.ckpt-1000.meta')
saver.restore(sess, tf.train.latest_checkpoint('./'))
此方法加载指定文件夹下最近保存的一个模型的数据；或者
with tf.Session() as sess:
saver = tf.train.import_meta_graph('model.ckpt-1000.meta')
saver.restore(sess, os.path.join(path, 'model.ckpt-1000'))
此方法可以指定具体某个数据，需要注意的是，指定的文件不要包含后缀。
使用Restore的模型将模型数据加载进来之后，下一步就是利用加载的模型进行下一步的操作了。这可以根据不同需要以如下几种方式进行操作。
1.查看模型参数可以直接查看Restore进来的模型的参数，如下：
with tf.Session() as sess:
saver = tf.train.import_meta_graph('model.ckpt-1000.meta')
saver.restore(sess, tf.train.latest_checkpoint('./'))
tvs = [v for v in tf.trainable_variables()]
for v in tvs:
print(v.name)
print(sess.run(v))
with tf.Session() as sess:
saver = tf.train.import_meta_graph('model.ckpt-1000.meta')
saver.restore(sess, tf.train.latest_checkpoint('./'))
gv = [v for v in tf.global_variables()]
for v in gv:
print(v.name)
上面通过global_variables()获得的与前trainable_variables类似，只是多了一些非trainable的变量，比如定义时指定为trainable=False的变量，或Optimizer相关的变量。
下面则可以获得几乎所有的operations相关的tensor：
with tf.Session() as sess:
saver = tf.train.import_meta_graph('model.ckpt-1000.meta')
saver.restore(sess, tf.train.latest_checkpoint('./'))
ops = [o for o in sess.graph.get_operations()]
for o in ops:
print(o.name)
首先，上面的sess.graph.get_operations()可以换为tf.get_default_graph().get_operations()，二者区别无非是graph明确的时候可以直接使用前者，否则需要使用后者。
此种方法获得的tensor比较齐全，可以从中一窥模型全貌。不过，最方便的方法还是推荐使用tensorboard来查看，当然这需要你提前将sess.graph输出。
2.直接使用原始模型进行训练或测试（前传）这种操作比较简单，无非是找到原始模型的输入、输出即可。
只要搞清楚输入输出的tensor名字，即可直接使用TensorFlow中graph的get_tensor_by_name函数，建立输入输出的tensor：
with tf.get_default_graph() as graph:
data = graph.get_tensor_by_name('data:0')
output = graph.get_tensor_by_name('output:0')
如上，需要特别注意，get_tensor_by_name后面传入的参数，如果没有重复，需要在后面加上“:0”。
从模型中找到了输入输出之后，即可直接使用其继续train整个模型，或者将输入数据feed到模型里，并前传得到test输出了。
需要说明的是，有时候从一个graph里找到输入和输出tensor的名字并不容易，所以，在定义graph时，最好能给相应的tensor取上一个明显的名字，比如：
data = tf.placeholder(tf.float32, shape=shape, name='input_data')
preds = tf.nn.softmax(logits, name='output')
诸如此类。这样，就可以直接使用tf.get_tensor_by_name(‘input_data:0’)之类的来找到输入输出了。
3.扩展原始模型除了直接使用原始模型，还可以在原始模型上进行扩展，比如对1中的output继续进行处理，添加新的操作，可以完成对原始模型的扩展，如：
with tf.get_default_graph() as graph:
data = graph.get_tensor_by_name('data:0')
output = graph.get_tensor_by_name('output:0')
logits = tf.nn.softmax(output)
4.使用原始模型的某部分有时候，我们有对某模型的一部分进行fine-tune的需求，比如使用一个VGG的前面提取特征的部分，而微调其全连层，或者将其全连层更换为使用convolution来完成，等等。TensorFlow也提供了这种支持，可以使用TensorFlow的stop_gradient函数，将模型的一部分进行冻结。
with tf.get_default_graph() as graph:
graph.get_tensor_by_name('fc1:0')
fc1 = tf.stop_gradient(fc1)
# add new procedure on fc1
| 
楼主
标题置顶
标题高亮
