Using disable_eager_execution also disables overriding train_step of model?

bach · October 6, 2022, 2:48pm

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")

chunduriv · October 7, 2022, 9:37am

This is due to

github.com

keras-team/keras/blob/master/keras/utils/version_utils.py#L40

      
        
            )
            callbacks = LazyLoader("callbacks", globals(), "keras.callbacks")
            callbacks_v1 = LazyLoader("callbacks_v1", globals(), "keras.callbacks_v1")
            
            

            
class ModelVersionSelector:
                """Chooses between Keras v1 and v2 Model class."""
            
            
    def __new__(cls, *args, **kwargs):
                    use_v2 = should_use_v2()
                    cls = swap_class(cls, training.Model, training_v1.Model, use_v2)
                    return super(ModelVersionSelector, cls).__new__(cls)
            
            

            
class LayerVersionSelector:
                """Chooses between Keras v1 and v2 Layer class."""
            
            
    def __new__(cls, *args, **kwargs):
                    use_v2 = should_use_v2()
                    cls = swap_class(cls, base_layer.Layer, base_layer_v1.Layer, use_v2)
                    return super(LayerVersionSelector, cls).__new__(cls)

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