Keras and AutoGraph

Viewed 457

Reading this and this answer I understood that with non-dynamical models on TensorFlow-2.0, Keras would use AutoGraph. But now writing a callback to get the variables' history during training,

class TrainHistory(tf.keras.callbacks.Callback):
    def on_train_begin(self, logs={}):
        self.vars = []

    def on_batch_end(self, batch, logs={}):
        self.vars.append([v.numpy() for v in self.model.variables])

I see that the .numpy() method of eager tensors can be used. I was expecting an error like numpy() is only available when eager execution is enabled. Is one executing eager code when using Keras with TensorFlow-2.0?

Here it's mentioned that for compute-intensive functions like convolutions there's not much speed up of tf.function-decorated functions compared to eager code, but from the numbers shown in the example I guess that the difference could be significative for long trainings. Would it be better performance-wise to use GradientTape on a custom training function decorated with tf.function instead of Keras' fit() method?

1 Answers

If you mean to use @tf.function on the outer loop, i.e. the 'epoch' loop, it would probably not benefit your model much. It would just make development harder. With more code comes complexity.

However, you must absolutely use tf.function on custom loss functions and other functions that get called one every batch.

And, no, using GradientTape-d custom tf.function would probably not out-perform the keras' fit method. It was thoroughly tested and refined over the years.

Does this answer your questions?

Related