In `tf.estimator`, how to `tf.assign` a variable at end of training (not at every iteration)?

Question

I am using the tf.estimator API to train models.

As I understand, the model_fn defines the computation graph, which returns a different tf.estimator.EstimatorSpec according to the mode.

In mode==tf.estimator.ModeKeys.TRAIN, one can specify a train_op to be called at each training iteration, which in turns changes trainable instances of tf.Variable, to optimise a certain loss.

Let's call the train_op optimizer, and the variables A and B.

In order to speed up prediction and evaluation, I would like to have an auxiliary non-trainable tf.Variable Tensor C, exclusively dependent on the already trained variables. The values of this tensor would thus be exportable. This Tensor does not affect training loss. Let's assume we want:

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

What I tried:

Passing tf.group(optimizer, update_op) as train_op in the EstimatorSpec works good but slows down training a lot, since the train_op now updates C at each iteration.

Because C is only needed at eval/predict time, one call of update_op at the end of training is enough.

Is it possible to assign a Variable at the end of training a tf.estimator.Estimator?

Why do you pass it into the estimator then? Just run update_op once at the end of the training manually. — Thomas Pinetz
The computation graph is defined in the model_fn, so it is update_op. The model_fn is passed as a parameter to the init method of a tf.estimator.Estimator. I do not see any trivial way running the update_op after call to estimator.train has finished. As a reminder, I'm putting the (reference)[tensorflow.org/get_started/custom_estimators] on the Estimator API @ThomasPinetz — syltruong

xdurch0 xdurch0 · Accepted Answer · 2018-05-18T12:03:55

In general a single iteration of the model function is not aware of whether training is going to end after it runs, so I doubt that this can be done straightforwardly. I see two options:

If you need the auxiliary variable only after training, you could use tf.estimator.Estimator.get_variable_value (see here) to extract the values of variables A and B after training as numpy arrays and do your computations to get C. However then C would not be part of the model.
Use a hook (see here). You can write a hook with an end method that will be called at the end of the session (i.e. when training stops). You would probably need to look into how hooks are defined/used -- e.g. here you can find implementations of most "basic" hooks already in Tensorflow. A rough skeleton could look something like this:
```
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)
```
Since the hooks needs access to the variables, you would need to define the hook inside the model function. You can pass such hooks to the training process in the EstimatorSpec (see here).

I haven't tested this! I'm not sure if you can define ops inside a hook. If not, it should hopefully work to define the update op inside the model function and pass that to the hook directly.

In `tf.estimator`, how to `tf.assign` a variable at end of training (not at every iteration)?

2 Answers