TensorFlow Estimator续训时如何恢复accuracy/total历史取值

根因 #

tf.compat.v1.metrics.accuracy 生成的 accuracy/total（正确预测累计数）、accuracy/count（总样本累计数）属于局部变量（local variable），并非全局变量。TensorFlow Estimator 默认的 Checkpoint 持久化逻辑仅保存 tf.GraphKeys.GLOBAL_VARIABLES 集合内的变量，不会主动留存局部变量。
因此续训时，Checkpoint中不存在这两个变量的历史值，TF会按照初始化规则将其重置为0，累计计数从头开始计算，直接导致精度随epoch的曲线出现从0起步的突跳。

修复实现 #

要在续训时完整恢复累计的total、count值，核心是将这两个变量纳入Checkpoint的保存范围，具体操作：

• 注册累计变量到可持久化集合
  在完成accuracy的图构建后，显式获取两个累计变量，手动加入全局可保存集合，让Estimator的默认Checkpoint逻辑自动将其写入ckpt文件：

  accuracy = tf.compat.v1.metrics.accuracy(
              labels=tf.argmax(labels, axis=1),
              predictions=tf.argmax(p_logits, axis=1),
  )
  # 获取metric对应的两个累计局部变量
  acc_total = tf.get_default_graph().get_tensor_by_name("accuracy/total:0")
  acc_count = tf.get_default_graph().get_tensor_by_name("accuracy/count:0")
  # 纳入全局可持久化集合，自动随Checkpoint保存
  tf.add_to_collection(tf.GraphKeys.GLOBAL_VARIABLES, acc_total)
  tf.add_to_collection(tf.GraphKeys.GLOBAL_VARIABLES, acc_count)
  

• 修正日志打印的变量映射
  tf.compat.v1.metrics.accuracy 返回的元组顺序为 (当前累计精度值, 累计值更新操作)，修正原有日志hook的取值，直接打印需要的累计值做监控：

  logging_hook = tf.estimator.LoggingTensorHook(
      {"acc_total": acc_total, "acc_count": acc_count, "cumulative_acc": accuracy[0]},
      every_n_iter=500
  )
  

• 可选的自定义Saver实现
  如果不想修改原有全局变量集合，可自定义Saver指定需要保存的变量列表，替换Estimator默认的Checkpoint保存逻辑：

  saver = tf.train.Saver(
      var_list=tf.global_variables() + [acc_total, acc_count]
  )
  # 替换默认Checkpoint Hook
  saver_hook = tf.estimator.CheckpointSaverHook(
      checkpoint_dir=model_dir,
      save_steps=1000,
      saver=saver
  )
  

注意事项 #

• 不要批量将所有局部变量加入全局可保存集合，仅注册需要持久化的metric累计变量即可，避免Checkpoint文件体积无意义膨胀
• 重新启动训练任务时，无需对acc_total、acc_count做额外的初始化操作，Estimator会自动从已有的Checkpoint中加载这两个变量的历史值，实现累计值的无缝衔接

内容的提问来源于stack exchange，提问作者Oscar Moreno