python - 在 `tf.estimator` 中，如何在训练结束时(不是每次迭代时)将变量设置为 `tf.assign` ？

Question

我正在使用 tf.estimator API 来训练模型。

据我了解，model_fn定义了计算图，它根据模式返回不同的tf.estimator.EstimatorSpec。

在 mode==tf.estimator.ModeKeys.TRAIN 中，可以指定在每次训练迭代时调用的 train_op，这又会改变 trainable tf.Variable 实例，以优化一定的损失。

让我们调用 train_op 优化器，以及变量 A 和 B。

为了加快预测和评估速度，我想要一个辅助的不可训练 tf.Variable 张量C，完全依赖于已经训练过的变量。因此，该张量的值是可导出的。该张量不影响训练损失。假设我们想要:

C = tf.Variable(tf.matmul(A,B))
update_op = tf.assign(C, tf.matmul(A,B))

• 我尝试过的:

将 tf.group(optimizer, update_op) 作为 EstimatorSpec 中的 train_op 传递效果很好，但会大大减慢训练速度，因为 train_op 现在在每次迭代时更新 C。

由于仅在评估/预测时需要 C，因此在训练结束时调用一次 update_op 就足够了。

是否可以在训练tf.estimator.Estimator结束时分配一个变量？

Answer 1

一般来说，模型函数的单次迭代并不知道训练在运行后是否会结束，所以我怀疑这是否可以直接完成。我看到两个选项:

1. 如果您仅在训练后需要辅助变量，则可以使用 tf.estimator.Estimator.get_variable_value (参见 here )来提取变量 A< 的值 和 B 在训练为 numpy 数组后进行计算以获得 C。但是，C 将不会成为模型的一部分。

2. 使用钩子(Hook)(参见 here )。您可以编写一个带有 end 方法的钩子(Hook)，该方法将在 session 结束时(即训练停止时)被调用。您可能需要研究如何定义/使用钩子(Hook)——例如here您可以在 Tensorflow 中找到大多数“基本”钩子(Hook)的实现。粗略的骨架可能看起来像这样:

   class UpdateHook(SessionRunHook):
       def __init__(update_variable, other_variables):
           self.update_op = tf.assign(update_variable, some_fn(other_variables))
   
       def end(session):
           session.run(self.update_op)
   

   由于钩子(Hook)需要访问变量，因此您需要在模型函数内定义钩子(Hook)。您可以将此类 Hook 传递到 EstimatorSpec 中的训练过程(请参阅 here )。

   我还没有测试过这个!我不确定您是否可以在钩子(Hook)内定义操作。如果没有，它应该可以在模型函数内定义更新操作并将其直接传递给钩子(Hook)。