如何在estimator.train的每个步骤之间运行张量流操作

Question

我正在使用估算器来训练我的模型。

图中有一些操作会更新一些不可训练的变量。

如何在estimator.train函数中使用钩子来运行这些操作

Answer 1

from tensorflow.python.training import training

class MySessionRunHook(training.SessionRunHook):
  # your hook definition here

  def __init__():
    # ...

  def after_run(self, run_context, run_values):
    # Get your ops
    your_ops = ...
    run_context.session.run(your_ops)

hook = MySessionRunHook(...)
estimator.train(train_input_fn, hooks=[hook])