基于TensorFlow的slim模块的模型配置（DeepLab源程序分析）

在之前一篇博客中，借助DeepLabv3+的程序详细讲解了PASCAL-VOC2012数据集的读入过程，今天重点讲一下DeepLab中的训练过程模型配置与tensorboard可视化操作。DeepLab的代码模块结构非常清晰，便于对每一部分进行单独分析，而且DeepLab模型是当前图像语义分割领域最好的模型之一，所以我以该模型代码为例，详细叙述它的模型配置与tensorboard可视化操作过程。之前的博客TensorFlow数据集操作（使用slim和tfrecord）已经将数据集整理成为batch队列的形式，接下来需要对它的训练过程进行配置。

在该程序中对整个训练过程的配置主要使用了slim库中的deployment中的model_deploy文件。该文件中定义了一个DeploymentConfig类以及众多与训练过程配置有关的类函数。首先看一下该文件中的使用方法的介绍。可以看到整个方法结构非常清晰层次明确，将整个模型训练过程分为几个小的独立的模块，包括创建config对象，配置训练时的变量，配置输入数据，配置优化方法，定义网络模型以及最后的模型训练。

Usage:
  g = tf.Graph() # 定义计算图
  # 创建DeploymentConfig类对象config
  config = model_deploy.DeploymentConfig(num_clones=2, clone_on_cpu=True)
  # 配置并创建global step来存储变量
  with tf.device(config.variables_device()): 
    global_step = slim.create_global_step() 
  # 配置输入数据
  with tf.device(config.inputs_device()): 
    images, labels = LoadData(...)
    inputs_queue = slim.data.prefetch_queue((images, labels))
  # 配置优化方法
  with tf.device(config.optimizer_device()): 
    optimizer = tf.train.MomentumOptimizer(FLAGS.learning_rate, FLAGS.momentum)
  # 定义网络模型
  def model_fn(inputs_queue):
    images, labels = inputs_queue.dequeue() # 从队列中出列图像标签数据
    predictions = CreateNetwork(images) # 神经网络前向传播
    slim.losses.log_loss(predictions, labels) # 定义损失
  # 根据网络，输入数据，优化器定义训练模型
  model_dp = model_deploy.deploy(config, model_fn, [inputs_queue],
                                 optimizer=optimizer)
  # 运行训练过程
  slim.learning.train(model_dp.train_op, my_log_dir,
                      summary_op=model_dp.summary_op)

下面来具体的看一下DeepLab程序中是如何使用的。可以看到该程序主要都是以config对象对整个训练过程进行配置，这正是slim中定义的DeploymentConfig类的作用。整个过程主要包括三步：（1）定义配置类对象config用于配置训练过程；（2）在TensorFlow计算图中对各个部分进行配置，包括输入数据，变量操作，神经网络模型，优化过程。（3）创建计算图运行变量Session，并通过slim.learning.train函数运行计算图，开始训练过程。本文中程序的主要难点在于设置多GPU运行，需要考虑一些相关参数。该程序的模型具体配置过程与其他程序几乎没有区别，都是采用相同的模型配置方法。

注：这段代码仅仅是源DeepLab代码的部分简化版，去掉了不相关的函数调用与模型可视化部分，仅保留模型配置代码。另外形如FLAGS.***的参数都是运行程序时从命令行得到的参数。

tf.logging.set_verbosity(tf.logging.INFO) # 显示INFO及更高级别的日志消息

# 创建DeploymentConfig对象config,这个配置类描述了如何将一个模型部署在
# 多个单机的多个GPU上，在每个单机上，模型将被复制num_clones次
config = model_deploy.DeploymentConfig(
      num_clones=FLAGS.num_clones,       # 每个单机部署多少个clone（即部署在多少个GPU）
      clone_on_cpu=FLAGS.clone_on_cpu,   # 如果为True，则单机中的每个clone将被放在CPU中
      replica_id=FLAGS.task,     		 # 整数，模型所部署的单机的索引，通常是0
      num_replicas=FLAGS.num_replicas,   # 使用多少个单机，通常为1，表示单机部署。此时`worker_device`, `num_ps_tasks`和 `ps_device`这几个参数将被忽略。
      num_ps_tasks=FLAGS.num_ps_tasks)   # 分布式作业使用多少个单机，如果为0表示不使用单机
# 一个btch的数据要被同时平均分配给所有的GPU，此时每个GPU的batch_size
clone_batch_size = FLAGS.train_batch_size // config.num_clones

# 定义网络模型
def _build_deeplab(inputs_queue, outputs_to_num_classes, ignore_label):
	... ...
	return outputs_to_scales_to_logits

with tf.Graph().as_default() as graph:
	# 配置模型变量
	with tf.device(config.variables_device()):
		global_step = tf.train.get_or_create_global_step() # 创建global-step，参数为计算图graph，如果没有参数则采用默认计算图
		# Define the model and create clones.
		model_fn = _build_deeplab # 引用模型，模型的定义在下面具体的函数中
		model_args = (inputs_queue, {common.OUTPUT_TYPE: dataset.num_classes}, dataset.ignore_label)
		clones = model_deploy.create_clones(config, model_fn, args=model_args) # 创建模型并克隆到多个GPU	
	  	update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS, first_clone_scope) # 返回计算图中的first_clone_scope空间中的名字为UPDATE_OPS的张量集合（即返回需要迭代更新的变量）

	# 配置模型的输入数据	
	with tf.device(config.inputs_device()): 
		inputs_queue = prefetch_queue.prefetch_queue(samples, capacity=128 * config.num_clones) # 将batch数据预存入队列
		first_clone_scope = config.clone_scope(0) # 根据索引号0返回创建的克隆的scope名字first_clone_scope
    
    # 配置优化过程
	with tf.device(config.optimizer_device()):
		learning_rate = train_utils.get_model_learning_rate(  # 设定学习率（poly）
        	FLAGS.learning_policy, FLAGS.base_learning_rate,
        	FLAGS.learning_rate_decay_step, FLAGS.learning_rate_decay_factor,
        	FLAGS.training_number_of_steps, FLAGS.learning_power,
        	FLAGS.slow_start_step, FLAGS.slow_start_learning_rate)
		optimizer = tf.train.MomentumOptimizer(learning_rate, FLAGS.momentum) # 设定优化方法（动量随机梯度下降）
		startup_delay_steps = FLAGS.task * FLAGS.startup_delay_steps # 这一条我没有看懂
	
	# 配置模型变量
	with tf.device(config.variables_device()):
		total_loss, grads_and_vars = model_deploy.optimize_clones(clones, optimizer) # 根据clones网络模型和optimize优化方法计算损失和梯度
		total_loss = tf.check_numerics(total_loss, 'Loss is inf or nan.') # 检查检查 NaN 和 Inf 值的张量，如果存在这些张量则返回异常信息

		# 修正bias和最后一层权值变量的梯度。对于语义分割任务，模型通常是从分类任务中微调而来。
		# 为了微调模型，我们通常把模型的最后一层变量设置更大的学习率
		last_layers = model.get_extra_layer_scopes(FLAGS.last_layers_contain_logits_only)
		grad_mult = train_utils.get_model_gradient_multipliers(last_layers, FLAGS.last_layer_gradient_multiplier)
		if grad_mult:
			grads_and_vars = slim.learning.multiply_gradients(grads_and_vars, grad_mult)

		# 创建梯度更新对象
		grad_updates = optimizer.apply_gradients(grads_and_vars, global_step=global_step)
		update_ops.append(grad_updates)
		update_op = tf.group(*update_ops)
		# 用来控制计算流图，给图中的某些计算指定顺序，即先运行update_op再运行total_loss，即更新完变量后再计算损失
		# 与tf.control_dependencies配套使用的tf.identity用于创建一个与原来一样的张量节点到graph中，这样control_dependencies才会生效
		with tf.control_dependencies([update_op]):
			train_tensor = tf.identity(total_loss, name='train_op')
	
	# 创建session，并对session进行参数配置，指定在GPU设备上的运行情况
	session_config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)

    # 运行训练过程
    slim.learning.train(
        train_tensor,
        logdir=FLAGS.train_logdir,
        log_every_n_steps=FLAGS.log_steps,
        master=FLAGS.master,
        number_of_steps=FLAGS.training_number_of_steps,
        is_chief=(FLAGS.task == 0),
        session_config=session_config,
        startup_delay_steps=startup_delay_steps,
        init_fn=train_utils.get_model_init_fn(
            FLAGS.train_logdir,
            FLAGS.tf_initial_checkpoint,
            FLAGS.initialize_last_layer,
            last_layers,
            ignore_missing_vars=True),
        summary_op=summary_op,
        save_summaries_secs=FLAGS.save_summaries_secs,
        save_interval_secs=FLAGS.save_interval_secs)