Tensorflow加载预训练模型和保存模型（ckpt文件）以及迁移学习finetuning

待我称王封你为后i 2022-09-04 11:00 536阅读 0赞

使用tensorflow过程中，训练结束后我们需要用到模型文件。有时候，我们可能也需要用到别人训练好的模型，并在这个基础上再次训练。这时候我们需要掌握如何操作这些模型数据。看完本文，相信你一定会有收获！

# 1 Tensorflow模型文件 #

我们在checkpoint\_dir目录下保存的文件结构如下：

## 1.1 meta文件 ##

MyModel.meta文件保存的是图结构，meta文件是pb（protocol buffer）格式文件，包含变量、op、集合等。

## 1.2 ckpt文件 ##

ckpt文件是二进制文件，保存了所有的weights、biases、gradients等变量。在tensorflow 0.11之前，保存在**.ckpt**文件中。0.11后，通过两个文件保存,如：

1.  `MyModel.data-00000-of-00001`
2.  `MyModel.index`

## 1.3 checkpoint文件 ##

我们还可以看，checkpoint\_dir目录下还有checkpoint文件，该文件是个文本文件，里面记录了保存的最新的checkpoint文件以及其它checkpoint文件列表。在inference时，可以通过修改这个文件，指定使用哪个model

# 2 保存Tensorflow模型 #

tensorflow 提供了`tf.train.Saver`类来保存模型，值得注意的是，在tensorflow中，变量是存在于Session环境中，也就是说，只有在Session环境下才会存有**变量**值，因此，保存模型时需要传入session：

1.  `saver = tf.train.Saver()`
2.  `saver.save(sess,"./checkpoint_dir/MyModel")`

看一个简单例子：

1.  `import tensorflow as tf`
2.  
3.  `w1 = tf.Variable(tf.random_normal(shape=[2]), name='w1')`
4.  `w2 = tf.Variable(tf.random_normal(shape=[5]), name='w2')`
5.  `saver = tf.train.Saver()`
6.  `sess = tf.Session()`
7.  `sess.run(tf.global_variables_initializer())`
8.  `saver.save(sess, './checkpoint_dir/MyModel')`

执行后，在checkpoint\_dir目录下创建模型文件如下：

1.  `checkpoint`
2.  `MyModel.data-00000-of-00001`
3.  `MyModel.index`
4.  `MyModel.meta`

另外，如果想要在1000次迭代后，再保存模型，只需设置`global_step`参数即可

保存的模型文件名称会在后面加`-1000`,如下：

1.  `checkpoint`
2.  `MyModel-1000.data-00000-of-00001`
3.  `MyModel-1000.index`
4.  `MyModel-1000.meta`

在实际训练中，我们可能会在每1000次迭代中保存一次模型数据，但是由于图是不变的，没必要每次都去保存，可以通过如下方式指定不保存图：

saver.save(sess, './checkpoint_dir/MyModel',global_step=step,write_meta_graph=False)

另一种比较实用的是，如果你希望每2小时保存一次模型，并且只保存最近的5个模型文件：

tf.train.Saver(max_to_keep=5, keep_checkpoint_every_n_hours=2)

> 注意：tensorflow默认只会保存最近的5个模型文件，如果你希望保存更多，可以通过`max_to_keep`来指定

如果我们不对`tf.train.Saver`指定任何参数，默认会保存所有变量。如果你不想保存所有变量，而只保存一部分变量，可以通过指定variables/collections。在创建`tf.train.Saver`实例时，通过将需要保存的变量构造list或者dictionary，传入到Saver中：

1.  `import tensorflow as tf`
2.  `w1 = tf.Variable(tf.random_normal(shape=[2]), name='w1')`
3.  `w2 = tf.Variable(tf.random_normal(shape=[5]), name='w2')`
4.  `saver = tf.train.Saver([w1,w2])`
5.  `sess = tf.Session()`
6.  `sess.run(tf.global_variables_initializer())`
7.  `saver.save(sess, './checkpoint_dir/MyModel',global_step=1000)`

# 3 导入训练好的模型 #

在第1小节中我们介绍过，tensorflow将图和变量数据分开保存为不同的文件。因此，在导入模型时，也要分为2步：构造网络图和加载参数

## 3.1 构造网络图 ##

一个比较笨的方法是，手敲代码，实现跟模型一模一样的图结构。其实，我们既然已经保存了图，那就没必要在去手写一次图结构代码。

saver=tf.train.import_meta_graph('./checkpoint_dir/MyModel-1000.meta')

上面一行代码，就把图加载进来了

## 3.2 加载参数 ##

仅仅有图并没有用，更重要的是，我们需要前面训练好的模型参数（即weights、biases等），本文第2节提到过，变量值需要依赖于Session，因此在加载参数时，先要构造好Session：

1.  `import tensorflow as tf`
2.  `with tf.Session() as sess:`
3.  `new_saver = tf.train.import_meta_graph('./checkpoint_dir/MyModel-1000.meta')`
4.  `new_saver.restore(sess, tf.train.latest_checkpoint('./checkpoint_dir'))`

此时，W1和W2加载进了图，并且可以被访问：

1.  `import tensorflow as tf`
2.  `with tf.Session() as sess:`
3.  `saver = tf.train.import_meta_graph('./checkpoint_dir/MyModel-1000.meta')`
4.  `saver.restore(sess,tf.train.latest_checkpoint('./checkpoint_dir'))`
5.  `print(sess.run('w1:0'))`
6.  `##Model has been restored. Above statement will print the saved value`

执行后，打印如下：

[ 0.51480412 -0.56989086]

# 4 使用恢复的模型 #

前面我们理解了如何保存和恢复模型，很多时候，我们希望使用一些已经训练好的模型，如prediction、fine-tuning以及进一步训练等。这时候，我们可能需要获取训练好的模型中的一些中间结果值，可以通过`graph.get_tensor_by_name('w1:0')`来获取，注意`w1:0`是tensor的name。

假设我们有一个简单的网络模型，代码如下：

1.  `import tensorflow as tf`
2.  
3.  
4.  `w1 = tf.placeholder("float", name="w1")`
5.  `w2 = tf.placeholder("float", name="w2")`
6.  `b1= tf.Variable(2.0,name="bias")`
7.  
8.  `#定义一个op，用于后面恢复`
9.  `w3 = tf.add(w1,w2)`
10. `w4 = tf.multiply(w3,b1,name="op_to_restore")`
11. `sess = tf.Session()`
12. `sess.run(tf.global_variables_initializer())`
13. 
14. `#创建一个Saver对象，用于保存所有变量`
15. `saver = tf.train.Saver()`
16. 
17. `#通过传入数据，执行op`
18. `print(sess.run(w4,feed_dict ={w1:4,w2:8}))`
19. `#打印 24.0 ==>(w1+w2)*b1`
20. 
21. `#现在保存模型`
22. `saver.save(sess, './checkpoint_dir/MyModel',global_step=1000)`

接下来我们使用`graph.get_tensor_by_name()`方法来操纵这个保存的模型。

1.  `import tensorflow as tf`
2.  
3.  `sess=tf.Session()`
4.  `#先加载图和参数变量`
5.  `saver = tf.train.import_meta_graph('./checkpoint_dir/MyModel-1000.meta')`
6.  `saver.restore(sess, tf.train.latest_checkpoint('./checkpoint_dir'))`
7.  
8.  
9.  `# 访问placeholders变量，并且创建feed-dict来作为placeholders的新值`
10. `graph = tf.get_default_graph()`
11. `w1 = graph.get_tensor_by_name("w1:0")`
12. `w2 = graph.get_tensor_by_name("w2:0")`
13. `feed_dict ={w1:13.0,w2:17.0}`
14. 
15. `#接下来，访问你想要执行的op`
16. `op_to_restore = graph.get_tensor_by_name("op_to_restore:0")`
17. 
18. `print(sess.run(op_to_restore,feed_dict))`
19. `#打印结果为60.0==>(13+17)*2`

> 注意：保存模型时，只会保存变量的值，placeholder里面的值不会被保存

如果你不仅仅是用训练好的模型，还要加入一些op，或者说加入一些layers并训练新的模型，可以通过一个简单例子来看如何操作：

1.  `import tensorflow as tf`
2.  
3.  `sess = tf.Session()`
4.  `# 先加载图和变量`
5.  `saver = tf.train.import_meta_graph('my_test_model-1000.meta')`
6.  `saver.restore(sess, tf.train.latest_checkpoint('./'))`
7.  
8.  `# 访问placeholders变量，并且创建feed-dict来作为placeholders的新值`
9.  `graph = tf.get_default_graph()`
10. `w1 = graph.get_tensor_by_name("w1:0")`
11. `w2 = graph.get_tensor_by_name("w2:0")`
12. `feed_dict = {w1: 13.0, w2: 17.0}`
13. 
14. `#接下来，访问你想要执行的op`
15. `op_to_restore = graph.get_tensor_by_name("op_to_restore:0")`
16. 
17. `# 在当前图中能够加入op`
18. `add_on_op = tf.multiply(op_to_restore, 2)`
19. 
20. `print (sess.run(add_on_op, feed_dict))`
21. `# 打印120.0==>(13+17)*2*2`

如果只想恢复图的一部分，并且再加入其它的op用于fine-tuning。只需通过`graph.get_tensor_by_name()`方法获取需要的op，并且在此基础上建立图，看一个简单例子，假设我们需要在训练好的VGG网络使用图，并且修改最后一层，将输出改为2，用于fine-tuning新数据：

1.  `......`
2.  `......`
3.  `saver = tf.train.import_meta_graph('vgg.meta')`
4.  `# 访问图`
5.  `graph = tf.get_default_graph()`
6.  
7.  `#访问用于fine-tuning的output`
8.  `fc7= graph.get_tensor_by_name('fc7:0')`
9.  
10. `#如果你想修改最后一层梯度，需要如下`
11. `fc7 = tf.stop_gradient(fc7) # It's an identity function`
12. `fc7_shape= fc7.get_shape().as_list()`
13. 
14. `new_outputs=2`
15. `weights = tf.Variable(tf.truncated_normal([fc7_shape[3], num_outputs], stddev=0.05))`
16. `biases = tf.Variable(tf.constant(0.05, shape=[num_outputs]))`
17. `output = tf.matmul(fc7, weights) + biases`
18. `pred = tf.nn.softmax(output)`
19. 
20. `# Now, you run this with fine-tuning data in sess.run()`

*  # 五、查看模型的所有层的输入输出的tensor name![70][] #

1.  `import os`
2.  `import re`
3.  `import tensorflow as tf`
4.  `from tensorflow.python import pywrap_tensorflow`
5.  `model_exp = "20180402-114759"`
6.  
7.  `def get_model_filenames(model_dir):`
8.  `files = os.listdir(model_dir)`
9.  `meta_files = [s for s in files if s.endswith('.meta')]`
10. `if len(meta_files)==0:`
11. `raise load_modelValueError('No meta file found in the model directory (%s)' % model_dir)`
12. `elif len(meta_files)>1:`
13. `raise ValueError('There should not be more than one meta file in the model directory (%s)' % model_dir)`
14. `meta_file = meta_files[0]`
15. `ckpt = tf.train.get_checkpoint_state(model_dir) # 通过checkpoint文件找到模型文件名`
16. `if ckpt and ckpt.model_checkpoint_path:`
17. `# ckpt.model_checkpoint_path表示模型存储的位置，不需要提供模型的名字，它回去查看checkpoint文件`
18. `ckpt_file = os.path.basename(ckpt.model_checkpoint_path)`
19. `return meta_file, ckpt_file`
20. 
21. `meta_files = [s for s in files if '.ckpt' in s]`
22. `max_step = -1`
23. `for f in files:`
24. `step_str = re.match(r'(^model-[\w\- ]+.ckpt-(\d+))', f)`
25. `if step_str is not None and len(step_str.groups())>=2:`
26. `step = int(step_str.groups()[1])`
27. `if step > max_step:`
28. `max_step = step`
29. `ckpt_file = step_str.groups()[0]`
30. `return meta_file, ckpt_file`
31. 
32. 
33. `meta_file, ckpt_file = get_model_filenames(model_exp)`
34. 
35. `print('Metagraph file: %s' % meta_file)`
36. `print('Checkpoint file: %s' % ckpt_file)`
37. `reader = pywrap_tensorflow.NewCheckpointReader(os.path.join(model_exp, ckpt_file))`
38. `var_to_shape_map = reader.get_variable_to_shape_map()`
39. `for key in var_to_shape_map:`
40. `print("tensor_name: ", key)`
41. `# print(reader.get_tensor(key))`
42. 
43. `with tf.Session() as sess:`
44. `saver = tf.train.import_meta_graph(os.path.join(model_exp, meta_file))`
45. `saver.restore(tf.get_default_session(),`
46. `os.path.join(model_exp, ckpt_file))`
47. `print(tf.get_default_graph().get_tensor_by_name("Logits/weights:0"))`

# 六、tensorflow从已经训练好的模型中，恢复指定权重(构建新变量、网络)并继续训练(finetuning) #

该部分转载自：[https://blog.csdn.net/ying86615791/article/details/76215363][https_blog.csdn.net_ying86615791_article_details_76215363]

假如要保存或者恢复指定tensor，并且把保存的graph恢复(插入)到当前的graph中呢？

总的来说，目前我会的是两种方法，命名都是很关键！  
两种方式保存模型，  
1.保存所有tensor，即整张图的所有变量，  
2.只保存指定scope的变量  
两种方式恢复模型，  
1.导入模型的graph，用该graph的saver来restore变量  
2.在新的代码段中写好同样的模型（变量名称及scope的name要对应），用默认的graph的saver来restore指定scope的变量

**两种保存方式：**  
**1.保存整张图，所有变量**

1.  `...`
2.  `init = tf.global_variables_initializer()`
3.  `saver = tf.train.Saver()`
4.  `config = tf.ConfigProto()`
5.  `config.gpu_options.allow_growth=True`
6.  `with tf.Session(config=config) as sess:`
7.  `sess.run(init)`
8.  `...`
9.  `writer.add_graph(sess.graph)`
10. `...`
11. `saved_path = saver.save(sess,saver_path)`
12. `...`

**2.保存图中的部分变量**

1.  `...`
2.  `init = tf.global_variables_initializer()`
3.  `vgg_ref_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='vgg_feat_fc')#获取指定scope的tensor`
4.  `saver = tf.train.Saver(vgg_ref_vars)#初始化saver时，传入一个var_list的参数`
5.  `config = tf.ConfigProto()`
6.  `config.gpu_options.allow_growth=True`
7.  `with tf.Session(config=config) as sess:`
8.  `sess.run(init)`
9.  `...`
10. `writer.add_graph(sess.graph)`
11. `...`
12. `saved_path = saver.save(sess,saver_path)`
13. `...`

**两种恢复方式：  
1.导入graph来恢复**

1.  `...`
2.  `vgg_meta_path = params['vgg_meta_path'] # 后缀是'.ckpt.meta'的文件`
3.  `vgg_graph_weight = params['vgg_graph_weight'] # 后缀是'.ckpt'的文件，里面是各个tensor的值`
4.  `saver_vgg = tf.train.import_meta_graph(vgg_meta_path) # 导入graph到当前的默认graph中，返回导入graph的saver`
5.  `x_vgg_feat = tf.get_collection('inputs_vgg')[0] #placeholder, [None, 4096]，获取输入的placeholder`
6.  `feat_decode = tf.get_collection('feat_encode')[0] #[None, 1024]，获取要使用的tensor`
7.  `"""`
8.  `以上两个获取tensor的方式也可以为:`
9.  `graph = tf.get_default_graph()`
10. `centers = graph.get_tensor_by_name('loss/intra/center_loss/centers:0')`
11. `当然，前提是有tensor的名字`
12. `"""`
13. `...`
14. `init = tf.global_variables_initializer()`
15. `saver = tf.train.Saver() # 这个是当前新图的saver`
16. `config = tf.ConfigProto()`
17. `config.gpu_options.allow_growth=True`
18. `with tf.Session(config=config) as sess:`
19. `sess.run(init)`
20. `...`
21. `saver_vgg.restore(sess, vgg_graph_weight)#使用导入图的saver来恢复`
22. `...`

**2.重写一样的graph，然后恢复指定scope的变量**

1.  `def re_build():#重建保存的那个graph`
2.  `with tf.variable_scope('vgg_feat_fc'): #没错，这个scope要和需要恢复模型中的scope对应`
3.  `...`
4.  `return ...`
5.  
6.  `...`
7.  `vgg_ref_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='vgg_feat_fc') # 既然有这个scope，其实第1个方法中，导入graph后，可以不用返回的vgg_saver，再新建一个指定var_list的vgg_saver就好了，恩，需要传入一个var_list的参数`
8.  `...`
9.  `init = tf.global_variables_initializer()`
10. `saver_vgg = tf.train.Saver(vgg_ref_vars) # 这个是要恢复部分的saver`
11. `saver = tf.train.Saver() # 这个是当前新图的saver`
12. `config = tf.ConfigProto()`
13. `config.gpu_options.allow_growth=True`
14. `with tf.Session(config=config) as sess:`
15. `sess.run(init)`
16. `...`
17. `saver_vgg.restore(sess, vgg_graph_weight)#使用导入图的saver来恢复`
18. `...`

总结一下，这里的要点就是，在restore的时候，**saver要和模型对应**，如果直接用当前graph的saver = tf.train.Saver(),来恢复保存模型的权重saver.restore(vgg\_graph\_weight),就会报错，提示key/tensor ... not found之类的错误;  
写graph的时候，一定要注意**写好scope和tensor的name，合理插入variable\_scop**e;

最方便的方式还是，用第1种方式来保存模型，这样就不用重写代码段了，然后第1种方式恢复，不过为了稳妥，**最好还是通过获取var\_list，指定saver的var\_list**，妥妥的！

最新发现，用第1种方式恢复时，要记得当前的graph和保存的模型中**没有重名**的tensor，否则当前graph的tensor name可能不是那个name，可能在后面加了"\_1"....-\_-||

**在恢复图基础上构建新的网络（变量）并训练(finetuning)**

恢复模型graph和weights在上面已经说了，这里的关键点是怎样**只恢复原图的权重** ，并且使**optimizer只更新新构造变量（指定层、变量）**。

（以下code与上面没联系）

1.  `"""1.Get input, output , saver and graph"""#从导入图中获取需要的东西`
2.  `meta_path_restore = model_dir + '/model_'+model_version+'.ckpt.meta'`
3.  `model_path_restore = model_dir + '/model_'+model_version+'.ckpt'`
4.  `saver_restore = tf.train.import_meta_graph(meta_path_restore) #获取导入图的saver，便于后面恢复`
5.  `graph_restore = tf.get_default_graph() #此时默认图就是导入的图`
6.  `#从导入图中获取需要的tensor`
7.  `#1. 用collection来获取`
8.  `input_x = tf.get_collection('inputs')[0]`
9.  `input_is_training = tf.get_collection('is_training')[0]`
10. `output_feat_fused = tf.get_collection('feat_fused')[0]`
11. `#2. 用tensor的name来获取`
12. `input_y = graph_restore.get_tensor_by_name('label_exp:0')`
13. `print('Get tensors...')`
14. `print('inputs shape: {}'.format(input_x.get_shape().as_list()))`
15. `print('input_is_training shape: {}'.format(input_is_training.get_shape().as_list()))`
16. `print('output_feat_fused shape: {}'.format(output_feat_fused.get_shape().as_list()))`
17. 
18. 
19. `"""2.Build new variable for fine tuning"""#构造新的variables用于后面的finetuning`
20. `graph_restore.clear_collection('feat_fused') #删除以前的集合，假如finetuning后用新的代替原来的`
21. `graph_restore.clear_collection('prob')`
22. `#添加新的东西`
23. `if F_scale is not None and F_scale!=0:`
24. `print('F_scale is not None, value={}'.format(F_scale))`
25. `feat_fused = Net_normlize_scale(output_feat_fused, F_scale)`
26. `tf.add_to_collection('feat_fused',feat_fused)#重新添加到新集合`
27. `logits_fused = last_logits(feat_fused,input_is_training,7) # scope name是"final_logits"`
28. 
29. 
30. `"""3.Get acc and loss"""#构造损失`
31. `with tf.variable_scope('accuracy'):`
32. `accuracy,prediction = ...`
33. `with tf.variable_scope('loss'):`
34. `loss = ...`
35. 
36. `"""4.Build op for fine tuning"""`
37. `global_step = tf.Variable(0, trainable=False,name='global_step')`
38. `learning_rate = tf.train.exponential_decay(initial_lr,`
39. `global_step=global_step,`
40. `decay_steps=decay_steps,`
41. `staircase=True,`
42. `decay_rate=0.1)`
43. `update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)`
44. `with tf.control_dependencies(update_ops):`
45. `var_list = tf.contrib.framework.get_variables('final_logits')#关键！获取指定scope下的变量`
46. `train_op = tf.train.MomentumOptimizer(learning_rate=learning_rate,momentum=0.9).minimize(loss,global_step=global_step,var_list=var_list) #只更新指定的variables`
47. `"""5.Begin training"""`
48. `init = tf.global_variables_initializer()`
49. `saver = tf.train.Saver()`
50. `config = tf.ConfigProto()`
51. `config.gpu_options.allow_growth=True`
52. `with tf.Session(config=config) as sess:`
53. `sess.run(init)`
54. `saver_restore.restore(sess, model_path_restore) #这里saver_restore对应导入图的saver， 如果用上面新的saver的话会报错 因为多出了一些新的变量 在保存的模型中是没有那些权值的`
55. `sess.run(train_op, feed_dict)`
56. `.......`

**再说明下两个关键点：**

**1. 如何在新图的基础上 只恢复 导入图的权重 ？**

用导入图的saver: saver\_restore

**2. 如何只更新指定参数？**

用var\_list = tf.contrib.framework.get\_variables(scope\_name)获取指定scope\_name下的变量，

然后optimizer.minimize()时传入指定var\_list

**附：如何知道tensor名字以及获取指定变量？**

**1.获取某个操作之后的输出**

用**graph.get\_operations()**获取所有op

比如<tf.Operation 'common\_conv\_xxx\_net/common\_conv\_net/flatten/Reshape' type=Reshape>,

那么output\_pool\_flatten = graph\_restore.get\_tensor\_by\_name('common\_conv\_xxx\_net/common\_conv\_net/flatten/Reshape:0')就是那个位置经过flatten后的输出了

**2.获取指定的var的值**

用GraphKeys获取变量

**tf.get\_collection(tf.GraphKeys.TRAINABLE\_VARIABLES)**返回指定集合的变量

比如 <tf.Variable 'common\_conv\_xxx\_net/final\_logits/logits/biases:0' shape=(7,) dtype=float32\_ref>

那么var\_logits\_biases = graph\_restore.get\_tensor\_by\_name('common\_conv\_xxx\_net/final\_logits/logits/biases:0')就是那个位置的biases了

**3.获取指定scope的collection**

tf.get\_collection(tf.GraphKeys.REGULARIZATION\_LOSSES,scope='common\_conv\_xxx\_net.final\_logits')

注意后面的scope是xxx.xxx不是xxx/xxx

# Reference #

[http://cv-tricks.com/tensorflow-tutorial/save-restore-tensorflow-models-quick-complete-tutorial/][http_cv-tricks.com_tensorflow-tutorial_save-restore-tensorflow-models-quick-complete-tutorial]

[70]: /images/20220829/89dcb1da3e944ef895a2d4bdb01bc8e5.png
[https_blog.csdn.net_ying86615791_article_details_76215363]: https://blog.csdn.net/ying86615791/article/details/76215363
[http_cv-tricks.com_tensorflow-tutorial_save-restore-tensorflow-models-quick-complete-tutorial]: http://cv-tricks.com/tensorflow-tutorial/save-restore-tensorflow-models-quick-complete-tutorial/