Tensorflow框架 -- 模型的保存与恢复

一、模型保存与恢复

tensorflow 模型保存类，Class Saver()，初始化对象saver：

saver = tf.train.Saver()

def __init__(self,
             var_list=None,
             reshape=False,
             sharded=False,
             max_to_keep=5,
             keep_checkpoint_every_n_hours=10000.0,
             name=None,
             restore_sequentially=False,
             saver_def=None,
             builder=None,
             defer_build=False,
             allow_empty=False,
             write_version=saver_pb2.SaverDef.V2,
             pad_step_number=False,
             save_relative_paths=False,
             filename=None):

常用的参数如下：
  var_list：需要保存或恢复的的变量或者对象列表，默认是所有可保存的变量或者对象，通常使用默认值None。
  max_to_keep: 保存模型的数量，超过则会自动删除早期保存的模型。
  keep_checkpoint_every_n_hours：保存模型的间隔时间，默认情况下禁用该参数。

saver.save()

  def save(self,
           sess,
           save_path,
           global_step=None,
           latest_filename=None,
           meta_graph_suffix="meta",
           write_meta_graph=True,
           write_state=True,
           strip_default_attrs=False):

常用参数如下：
  sess：会话，保存模型必须要基于图和会话内。
  save_path：保存路径和名字。
  global_step：全局训练步骤，记录训练位置，可用于模型的恢复

saver.restore()

def restore(self, sess, save_path):
    Args:
      sess: A `Session` to use to restore the parameters.
      save_path: Path where parameters were previously saved.

在恢复预训练模型部分参数时，需要注意如下：

• 变量通常保存的形式为：<tf.Variable 'resnet_v2_50/conv1/weights:0' shape=(7, 7, 3, 64) dtype=float64>
• load = tf.train.Saver()初始化对象时var_list参数，那么在调用load.restore()时，恢复的权重就是var_list中名字对应的权重值。

二、Fine-tune下模型的保存与恢复

深度学习的任务都是在已经训练好的模型上进行微调或者迁移学习。比如语义分割，目标检测任务，通常在分类数据集ImageNet上训练的模型作为主干网络（VGG，ResNet 作为 BackBone），然后迁移到其它任务场景继续训练。

下面以yolo-v3的部分代码为实例进行说明：

代码参考：https://github.com/YunYang1994/tensorflow-yolov3

代码中预加载的权重是DarkNet-53在ImageNet和coco（80类）上预训练网络。我们需要使用预训练的网络作为起点，进行一类目标检测的任务，因此网络的最后三层有变化。（？，？，？，255）=》（？，？，？，18）。所以，预训练模型的权重无法全部加载到新的DarkNet-53网络（事实上，只是最后三层的变化 [‘conv_sbbox’, ‘conv_mbbox’, ‘conv_lbbox’]）。

... ...

with tf.name_scope("define_first_stage_train"):
	self.first_stage_trainable_var_list = []
	for var in tf.trainable_variables():
		var_name = var.op.name
		var_name_mess = str(var_name).split('/')
		
		# 只将最后三层的变量进行微调优化，网络其它部分参数冻结（frozen），不参与训练
		# 当最后三层优化差不多了，然后放开整个网络，整体优化
		if var_name_mess[0] in ['conv_sbbox', 'conv_mbbox', 'conv_lbbox']:
			self.first_stage_trainable_var_list.append(var)
	
	first_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate)
	
	# 设定优化器优化变量列表
	part_var_op = first_stage_optimizer.minimize(self.loss, var_list=self.first_stage_trainable_var_list)

with tf.name_scope("define_second_stage_train"):
	second_stage_trainable_var_list = tf.trainable_variables()
	second_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate)
	
	# 设定优化器优化变量列表
	train_op_all_variables = first_stage_optimizer.minimize(self.loss, var_list=self.second_stage_trainable_var_list )

with tf.name_scope('loader_and_saver'):
	# 由于新的任务，最后三层的shape有变化，所以只能加载其它权重参数
	variables_to_restore = []
	for v in self.global_variables():
		if v.name.split('/')[0] not in ['conv_sbbox', 'conv_mbbox', 'conv_lbbox']:
			variables_to_restore.append(v)
	
	# 默认val_list=None，此时需要修改为 =》需要restore的变量，所以需要特别指定
	self.loader = tf.train.Saver(variables_to_restore)
	self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=100)

def train(self):
	self.sess.run(tf.global_variabels_initializer())
	start_epoch = 1
	try:
		print('=> Restoring weights form: %s ...' % self.initial_weight)
		sefl.loader.restore(self.sess, self.initial_weight)
	except:
		print('=> % s does not exist !!!' % self.initial_weight)
		print('=> Not it starts to train YOLOV3 from scratch ...')
		self.first_stage_epochs = 0 
	
	for epoch in range(start_epoch, 1 + self.first_stage_epoch + self.second_stage_epoch):
		self.sess.run([], feed_dict={})
	
		self.saver.save(self.sess, save_path, global_step=epoch)	

加载预训练权重的步骤：

1. 获取预训练模型每一层变量 var.

   variables_to_restore = []
   for v in self.global_variables():
   	# 打印所有变量
   	print(v)
   	# 确定需要恢复权重的变量，放入列表
   	if v.name.split('/')[0] not in ['conv_sbbox', 'conv_mbbox', 'conv_lbbox']:
   		variables_to_restore.append(v)
   

2. 确定要加载的权重，将变量加入variables_to_restore.

3. 初始化权重恢复对象

   # 初始化loader的时候，指定了需要恢复权重的变量列表
   self.loader = tf.train.Saver(variables_to_restore)
   

4. 加载指定变量的预训练权重

   try:
   	print('=> Restoring weights form: %s ...' % self.initial_weight)
   	sefl.loader.restore(self.sess, self.initial_weight)
   except:
   	print('=> % s does not exist !!!' % self.initial_weight)
   	print('=> Not it starts to train YOLOV3 from scratch ...')
   	self.first_stage_epochs = 0