在“tf.estimator”中，如何在训练结束时（不是每次迭代）“tf.assign”一个变量？答案

【问题标题】：In `tf.estimator`, how to `tf.assign` a variable at end of training (not at every iteration)?在“tf.estimator”中，如何在训练结束时（不是每次迭代）“tf.assign”一个变量？
【发布时间】：2023-03-12 04:54:01
【问题描述】：

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

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

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

让我们调用 train_op optimizer，变量 A 和 B。

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

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

我尝试了什么：

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

因为C 仅在评估/预测时需要，所以在训练结束时调用update_op 就足够了。

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

【问题讨论】：

那你为什么要把它传给估计器呢？只需在训练结束时手动运行一次 update_op。
计算图在model_fn中定义，所以是update_op。 model_fn 作为参数传递给 tf.estimator.Estimator 的 init 方法。在对estimator.train 的调用完成后，我看不到任何运行update_op 的简单方法。提醒一下，我将（参考）[tensorflow.org/get_started/custom_estimators] 放在 Estimator API @ThomasPinetz

标签： python tensorflow machine-learning deep-learning

【解决方案1】：

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

C

使用钩子（参见here）。您可以使用 end 方法编写一个钩子，该方法将在会话结束时（即训练停止时）调用。您可能需要研究挂钩是如何定义/使用的——例如here 你可以在 Tensorflow 中找到大多数“基本”钩子的实现。粗略的骨架可能看起来像这样：
```
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)
```
由于钩子需要访问变量，您需要在模型函数中定义钩子。您可以在EstimatorSpec（参见here）中将此类钩子传递给训练过程。

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

【讨论】：

第一个解决方案并不令人满意，因为C 不再是模型中的张量。但是我尝试了第二种解决方案，它按预期工作！据我了解tf.train.SessionRunHook 是一个用于定义与会话相关的钩子的接口。创建UpdateHook 后，我将其作为training_chief_hooks 传递给EstimatorSpec。训练结束后，Variable 确实被赋值了。
第二个解决方案对我不起作用：即使在会话关闭时运行了 end 方法，它对后来导出的变量的值也没有影响。

【解决方案2】：

使用钩子是一种解决方案。但请注意，如果要更改变量的值，则不能在 end() 函数中进行更改，因为更改的结果无法存储在检查点文件中。例如，如果您更改 after_run 函数中的值，则结果将存储在检查点中。

【讨论】：