tensorflow之Supervisor

本文是我对官方文档Supervisor: Training Helper for Days-Long Trainings 的编译。

你可以直接run一个training的op若干次来训练你的tensorflow模型，同时当你训练结束的时候你可以保存训练参数的checkpoint。对于那些数小时便可训练好的小模型来说这样干是很不错的。
然而对于需要很多天训练的大模型来说，就需要更鲁棒的训练过程了。这要求训练时要:

• 干净的处理shutdown以及crash
• 在shutdown或者crash之后可以恢复
• 可以通过tensorboard来监控

为了在shutdown之后可以恢复训练，训练过程中必须规律性地run summary 的op，同时将返回的值加到事件文件（events file）中。 tensorboard监控事件文件并且显式图，来报告随着时间进行的训练过程。

tf.train.Supervisor 提供了一系列服务来帮助实现一个鲁棒的训练过程。

本文展示如何直接使用supervisor。

很简单的方案
最简单的使用supervisor的方案就是去：

• 创建一个Supervisor对象，将要保存checkpoints以及summaries的目录路径传递给该对象。
• 利用tf.train.Supervisor.managed_session向supervisor请求一个session。
• 利用该session来执行训练的op，在每一步都核查supervisor是否要求训练结束。

...create graph...
  my_train_op = ...

  sv = tf.train.Supervisor(logdir="/my/training/directory")
  with sv.managed_session() as sess:
    for step in range(100000):
      if sv.should_stop():
        break
      sess.run(my_train_op)

启动服务
在这个最简单的情景中，managed_session()调用开启了一些列服务，这些服务在它们各自的线程中运作，同时利用该托管的session在你的图中run ops。

如果你的图中有一个name为global_step的整型变量，服务会使用它的值来衡量执行的训练步数。

• checkpointing 服务：在logdir里保存图中变量的一份拷贝。如果图中添加了global_step变量，则checkpoint文件名使用该变量的值。默认每十分钟保存一次。
• summary服务：run所有的summary ops，同时将它们的输出加到logdir中的事件文件中。默认为两分钟执行一次。
• 步数（step）计数：计算已经有多少步被执行了，从变量global_step的变化中来得到。向事件文件中添加一个summary用来报告每秒global steps的数目。summary标签为”global_step/sec”。同样野默认为两分钟执行一次。
• Queue Runners,如果有图中添加有tf.train.QueueRunner，superivsor会在自己的线程中开始它们。

所有的时间间隔都可以在创建supervisor对象的时候改变。

对于停止的核查

训练的主循环中队与stop的check时重要且必须的。
当shold_stop()条件设置为true时，这些服务线程中提起的异常会被报告给supervisor。服务线程会通知该条件并且恰当地终止。在managed_session()块中的主循环，同样也一定会checkstop的condition同时终止。

注意，managed_session()专注于捕获从训练循环中引起的异常并将它们报告给supervisor。主循环并不需要队异常做任何其他特殊的事情。只需要check stop condition。

恢复
如果训练程序shut down或者crash了，那么它最近的checkpoint以及事件文件就被留到了logdir里。当你重启程序的时候，managed_session会从最近的checkpoint中加载图，并且从中止的地方恢复训练。

一个新的事件文件会生成。当你开启tensorboard并且只想logdir时，它将直到如何来merge连个事件文件中的内容，并且将展示恢复于来自checkpoint中最后一个global step的训练。

大模型的情景
最简单的情景已经足以处理大多数小到中模型的训练。更大的模型也许会在运行summary sevice的时候耗尽内存：summary ops是与main loop中的train op一起并行地run的。这会导致内存使用达到通常使用的两倍多。

对于打得模型你可以通知supervisor不要运行summary服务，作为替代，你在自己的主训练循环中来运行：创建supervisor的时候传递summary_op=None。

以下代码在训练循环中，每100步运行一次summary op：

...create graph...
  my_train_op = ...
  my_summary_op = tf.summary.merge_all()

  sv = tf.train.Supervisor(logdir="/my/training/directory",
                     summary_op=None) # Do not run the summary service
  with sv.managed_session() as sess:
    for step in range(100000):
      if sv.should_stop():
        break
      if step % 100 == 0:
        _, summ = session.run([my_train_op, my_summary_op])
        sv.summary_computed(sess, summ)
      else:
        session.run(my_train_op)

预训练的模型情景
managed_session()调用很关心在session中初始化模型。模型会在可能的时候从一个checkpoint中加载，亦或从scratch中初始化。

一个常见的情景是要用加载的预训练的checkpoint来初始化模型，而该预训练模型和当前模型有些许的不同。

你可以通过gei supervisor传递”init function”的方式来加载预训练的checkpoint。这个函数只有在模型需要从scratch初始化时才被调用，而模型从logdir中的checkpoint恢复的时候并不会。

为了加载预训练模型，init 函数需要一个tf.train.Saver对象，所以你应该创建一个saver。新模型也许包含一些预训练的checkpoint中不存在的变量，所以这是一个很好的思想：这个saver必须只加载预训练的变量。如果你正在使用默认的saver，你会在尝试加载所有变量的时候得到一个错误。

 ...create graph...
  # Create a saver that restores only the pre-trained variables.
  pre_train_saver = tf.train.Saver([pre_train_var1, pre_train_var2])

  # Define an init function that loads the pretrained checkpoint.
  def load_pretrain(sess):
    pre_train_saver.restore(sess, "<path to pre-trained-checkpoint>")

  # Pass the init function to the supervisor.
  #
  # The init function is called _after_ the variables have been initialized
  # by running the init_op.
  sv = tf.train.Supervisor(logdir="/my/training/directory",
                     init_fn=load_pretrain)
  with sv.managed_session() as sess:
    # Here sess was either initialized from the pre-trained-checkpoint or
    # recovered from a checkpoint saved in a previous run of this code.

运行你自己的服务

Supervisor服务，比如checkpointing服务，与主训练循环并行运行。有时候你想加入你自己的服务，比如取出和 通常的summary的schedule不一样的不同设置的summaries。

使用supervisor中的tf.train.Supervisor.loop来达成这个目的。它会根据你选择的定时器重复地调用一个函数，直到supervisor的stop condition为true，所以它和其他服务工地很协调。

例如：每20分钟调用一次my_additional_summaries():

def my_additional_sumaries(sv, sess):
 ...fetch and write summaries, see below...

...
  sv = tf.train.Supervisor(logdir="/my/training/directory")
  with sv.managed_session() as sess:
    # Call my_additional_sumaries() every 1200s, or 20mn,
    # passing (sv, sess) as arguments.
    sv.loop(1200, my_additional_sumaries, args=(sv, sess))
    ...main training loop...

写summaries
supervisor总是在其logdir中生成一个事件文件，同时用一个tf.summary.FileWriter将事件和summaries添加到事件文件。如果你想哟啊写自己的summaries，也可以将它们添加到同一个事件文件中去：TensorBoard很喜欢在目录中只有一个事件文件。

supervisor提供了一个辅助函数来添加summaries：tf.train.Supervisor.summary_computed。仅仅需要将summary op的返回输出传递给该函数。以下是一个使用该函数实现之前例子中my_additional_sumaries()的例子：

def my_additional_sumaries(sv, sess):
  summaries = sess.run(my_additional_summary_op)
  sv.summary_computed(sess, summaries)

更多前沿的用法参看tf.train.Supervisor.summary_writer属性。

supervisor 引用
最简单的情景以及大模型的情景展示了supervisor的基本用法。更高级的情景可以用supervisor提供的很多选项来创建。

Checkpointing:何时何处
managed_session()调用开启了checkpointing服务，而这可以通过对Supervisor()创建时以下的参数来配置：

• logdir: checkpointing服务床及爱你checkpoints的目录路径。需要的化该目录会被创建。传递None禁用checkpointing以及summary服务。
• checkpoint_basename: 欲创建的checkpoint文件的名称，默认为”model.ckpt”。

如果模型包含一个标量整型变量名字为global_step，则该变量的值会被添加到checkpoint文件名中。

例如，在global step 1234，checkpoint 文件名就是 “model.ckpt-1234”。

• save_model_secs: 每个checkpoint之间的秒数。默认为600，即10分钟。
  当选择一个值时，要考虑一旦有crash时你要丢失多少工作：你永远不会丢失多于save_model_secs秒的工作。设置为0就禁用了checkpointing服务。
• saver： 一个tf.train.Saver对象，用来checkpointing。
  如果不传递saver，supervisor会调用tf.train.Saver()来创建一个，该saver会把所有的ops保存，并加载你模型中所有的变量。你通常也需要这么做。

例子：用一个习惯的Saver，每30秒checkpoint一次。

 ...create graph...
  my_saver = tf.train.Saver(<only some variables>)
  sv = tf.train.Supervisor(logdir="/my/training/directory",
                     saver=my_saver,
                     save_model_secs=30)
  with sv.managed_session() as sess:
    ...training loop...

Summaries:何时何处
类似checkpointing，logdir对summaries的作用也是一样的。事件文件在此创建，如果None则禁用了summary服务。

• save_summaries_secs：该参数代表每次运行summary sevice服务的间隔的秒数。默认为120秒，即两分钟。同样，设置为0时则禁用了summary服务。
• summary_op，用来取得summaries的op。
  如果没指定，supervisor会使用tf.GraphKeys.SUMMARY_OP图集合（graph collection）中第一个op。如果该集合为空，supervisor则创建一个op，它会将图中的所有summaries使用tf.summary.merge_all()聚集在一起。
  如果给summary_op传递None则禁用了summary服务。
• global_step：用来计算全局步数的张量。
  如果没有指明，supervisor使用tf.GraphKeys.GLOBAL_STEP图集合（graph collection）中第一个tensor，如果该集合为空，
  supervisor在图中寻找一个name为global_step的整型的变量的标量。
  如果找到，global step张量被用来衡量训练步数执行的数量。注意，你的训练op会增加global step的值。

模型的初始化和恢复
managed_session()调用野专注于初始化以及恢复一个session。它返回一个session同时伴随一个全部初始化了的模型，准备去训练。如果managed_session()调用时logdir里有一个checkpoint，模型会通过加载checkpoint初始化，否则会通过调用一个初始化op或者选择一个init function。
如果没有可用的checkpoint，模型的初始化则有下面的参数传递给supervisor()的创建器来控制：

• init_op： 需要被运行来初始化模型的op。
  如果没有指定，supervisor会使用tf.GraphKeys.INIT_OP图集合（ collection）中第一个op。如果集合是空的，则会通过调用tf.global_variables_initializer()添加一个初始化所有变量的op。
  传递None则不适用初始化op。

• init_fn: 调用它来初始化模型。
  如果指定则这样调用 ：init_fn(sess)，这里的sess是managed session。如果init op同时使用，则init function在init op之后被调用。

• local_init_op: 一个额外的op，用来初始化图段一部分，这部分没有被保存在checkpoints中。比如比如tables以及一些local variables。local init op在init op以及 init function之后运行。
  如果没有指定，supervisor使用tf.GraphKeys.LOCAL_INIT_OP集合里的第一个op。如果集合为空，则通过调用tf.tables_initializer() 和 tf.local_variables_initializer()添加一初始化所有tables以及local variables的op。
  传递None禁用local init op。

• ready_op: 核查模型是否被初始化的op。

运行了local init op，init op以及init function之后，supervisor会通过执行ready op来验证模型是否被完全初始化。如果初始化则该op返回空字符串，否则返回模型那部分未被初始化的一个描述。
如果未指定，supervisor会使用tf.GraphKeys.READY_OP 集合中的第一个op。若集合未空，supervisosr通过调用tf.report_uninitialized_variables()创建一个ready op来确保所有变量都被初始化。

传递None来禁用ready op。在这种情况下模型初始化之后不进行核查。
checkpoint的恢复是由以下传给superfisor()创建器的参数控制：

• logdir：寻找checkpoints的路径。checkpoint服务保存了一个metadata文件，名为 “checkpoint”，在这个checkpoint目录中指明最近的一个checkpoint的路径。
  这个文件是文本格式的。你可以手工编辑它来从一个不同于最近的checkpoint的checkpoint中恢复。
• ready_op：和上面的一样。ready op在加载checkpoint之前和之后运行。第一次运行检查模型是否需要被初始化，第二次验证模型完全被初始化。
• local_init_op：和上面的一样。local init op在第一次运行ready op之前运行，来初始化局部变量以及tables。
• saver：和上面的一样。用来加载checkpoint的的Saver对象。