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

Question

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

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

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

标签： tensorflow

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])