Using disable_eager_execution also disables overriding train_step of model?

TensorFlow General Discussion
, ,
1

Hi,

using Keras 2.10.0 rc3 (precompiled, on Ubuntu 22).

I noticed that if I use tf.compat.v1.disable_eager_execution(), then overriding a model train_step() does not work anymore.

I wonder whether this is a bug or an ‘expected behaviour’.

Code to reproduce:

import sys
import tensorflow as tf
import numpy as np
from tensorflow import keras

class CustomModel(keras.Sequential):
    def train_step(self, data):
        print("--- train step override")
        sys.exit(0)

# Commenting out the below line gives expected 
#  output ("--- train step override")
tf.compat.v1.disable_eager_execution()

model= CustomModel()
model.add(keras.layers.Dense(1, activation="sigmoid"))
model.compile(optimizer="adam", loss="mse", metrics=["mae"])

# simple test
x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=1)
print("I should not reach this point")
3

This is due to

https://github.com/keras-team/keras/blob/master/keras/utils/version_utils.py#L40

The training_v1 does not support custom train_step where as training.py is only used when executed eagerly. Thank you

1 Like