The problem with `tf.compat.v1.disable_v2_behavior`

[huang06]

Description

When importing some tensorflow_v2 modules, the subsequent TF2 functions cannot work properly.

I found out that some modules call the tensorflow.disable_v2_behavior function when they are imported. It will disable the subsequent TF2 behaviors.

cd submarine-sdk/pysubmarine/submarine

$ grep -rnH 'disable_v2_behavior'
ml/tensorflow_v2/layers/core.py:20:tf.compat.v1.disable_v2_behavior()
utils/tf_utils_v2.py:26:    tf.compat.v1.disable_v2_behavior()
utils/tf_utils_v2.py:57:        tf.compat.v1.disable_v2_behavior()

Furthermore, without the tensorflow.disable_v2_behavior, the tests still work as usual. Maybe the disable_v2_behavior is not necessary to the tf2 modules.

In summary, We should remove the tf.compat.v1.disable_v2_behavior statements to avoid the similar problem.

How to reproduce the TF2 error

Create Python venv.

cd submarine-sdk/pysubmarine
python3 -m venv .venv
source .venv/bin/activate
python3 -m pip install -e .[tf2]

Create demo.py and copy the following Python code into it.

# demo.py
import tensorflow as tf

from submarine.ml.tensorflow_v2.model import DeepFM   # call `tf.compat.v1.disable_v2_behavior` function


class LinearNNModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.dense1 = tf.keras.layers.Dense(1, activation=tf.nn.relu)  # One in and one out

    def call(self, x):
        y_pred = self.dense1(x)
        return y_pred


model = LinearNNModel()
input_arr = tf.random.uniform((1, 5))
model(input_arr)

model.save('sample-model')

Running the Python code. the error message shows that the TF2 function cannot work properly.

$ python3 demo.py

...
Traceback (most recent call last):
  File "/home/tom/huang06/submarine/submarine-sdk/pysubmarine/demo.py", line 21, in <module>
    model.save('sample-model')
  File "/home/tom/huang06/submarine/submarine-sdk/pysubmarine/.venv/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/tom/huang06/submarine/submarine-sdk/pysubmarine/.venv/lib/python3.10/site-packages/keras/saving/save.py", line 142, in save_model
    raise NotImplementedError(
NotImplementedError: Saving the model to HDF5 format requires the model to be a Functional model or a Sequential model. It does not work for subclassed models, because such models are defined via the body of a Python method, which isn't safely serializable. Consider saving to the Tensorflow SavedModel format (by setting save_format="tf") or using `save_weights`.

[cdmikechen]

@featherchen Hi~ There are some questions about the tf part of the codes that we need your help to answer, please. Hope you can help us when you have time. Thanks!