Repository: mlflow/mlflow

Issue:
[BUG] TF2/Keras autologging failed to log model with default h5 save format
### Willingness to contribute
The MLflow Community encourages bug fix contributions. Would you or another member of your organization be willing to contribute a fix for this bug to the MLflow code base?

- [ ] Yes. I can contribute a fix for this bug independently.
- [x] Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.
- [ ] No. I cannot contribute a bug fix at this time.

### System information
MLflow Version: 1.13.0
TensorFlow version: 2.3.1
Python Version: 3.8.5
Operating System: Mac OS X 10.15.6 19G2021

### Describe the problem

Keras/TF2 models can be saved in the `SavedModel` format from the MLflow version 1.13.0 when `saved_format='tf'` is provided to `mlflow.keras.log_model`. When `mlflow.tensorflow.autolog(log_models=True)` is used, some trained models can't be saved on the train end callback in the default 'h5' format.

It seems that this happens because of https://github.com/mlflow/mlflow/blob/cbe0627428bc2692d851f4180a2ea59073040a63/mlflow/tensorflow.py#L779-L781 This callback calls `mlflow.keras.log_model` with default save format, which is `h5`.

One way to address this is to update `mlflow.tensorflow.autolog` with the `save_format: str = 'h5'` parameter, and propagate `save_format` option to callbacks.

### Code to reproduce issue

```python
import mlflow
import mlflow.tensorflow
import tensorflow as tf
from tensorflow.keras.layers.experimental.preprocessing import TextVectorization

VOCAB_SIZE = 10
SEQUENCE_LENGTH = 16
EMBEDDING_DIM = 32

train_samples = ["this is an example", "another example"]
train_labels = [0.43, 0.23]

mlflow.tensorflow.autolog()

with mlflow.start_run():
    vectorizer_layer = TextVectorization(
        input_shape=(1,),
        max_tokens=VOCAB_SIZE,
        output_mode="int",
        output_sequence_length=SEQUENCE_LENGTH,
    )
    vectorizer_layer.adapt(train_samples)

    model = tf.keras.Sequential(
        [
            vectorizer_layer,
            tf.keras.layers.Embedding(
                VOCAB_SIZE,
                EMBEDDING_DIM,
                name="embedding",
                mask_zero=True,
                input_shape=(1,),
            ),
            tf.keras.layers.GlobalAveragePooling1D(),
            tf.keras.layers.Dense(32, activation="relu"),
            tf.keras.layers.Dense(1, activation="tanh"),
        ]
    )
    model.compile(optimizer="adam", loss="mse", metrics="mae")
    model.fit(train_samples, train_labels, epochs=1)
```

### Other info / logs

```
2020/12/25 16:15:23 INFO mlflow.utils.autologging_utils: tensorflow autologging will track hyperparameters, performance metrics, model artifacts, and lineage information for the current tensorflow workflow to the MLflow run with ID '334a468262d64e7392594a859d2e7da5'
2020-12-25 16:15:23.657399: I tensorflow/core/profiler/lib/profiler_session.cc:164] Profiler session started.
1/1 [==============================] - 0s 31ms/step - loss: 0.1191 - mae: 0.3325
2020/12/25 16:15:24 WARNING mlflow.utils.autologging_utils: MLflow issued a warning during tensorflow autologging: "/Users/user/.pyenv/versions/example/lib/python3.8/site-packages/mlflow/tensorflow.py:786: UserWarning: Logging to MLflow failed: Save or restore weights that is not an instance of `tf.Variable` is not supported in h5, use `save_format='tf'` instead. Got a model or layer TextVectorization with weights [<tensorflow.python.keras.engine.base_layer_utils.TrackableWeightHandler object at 0x14f40a460>]"
```

### What component(s), interfaces, languages, and integrations does this bug affect?

- [x] `area/tracking`: Tracking Service, tracking client APIs, autologging



Discussion:
has this issue  been fixed yet?

Not yet. Here is a quick workaround:

```python
_original_keras_log_model = mlflow.keras.log_model

def _keras_log_model_patch(*args, **kwargs) -> None:
    kwargs["save_format"] = "tf"
    _original_keras_log_model(*args, **kwargs)

mlflow.keras.log_model = _keras_log_model_patch
```

Resolving patch:
diff --git a/mlflow/keras.py b/mlflow/keras.py
--- a/mlflow/keras.py
+++ b/mlflow/keras.py
@@ -234,9 +234,8 @@ def _is_tf_keras(model):
     with open(os.path.join(data_path, _KERAS_MODULE_SPEC_PATH), "w") as f:
         f.write(keras_module.__name__)
 
-    # By default, Keras uses the SavedModel format -- specified by "tf"
-    # However, we choose to align with prior default of mlflow, HDF5
-    save_format = kwargs.get("save_format", "h5")
+    # Use the SavedModel format if `save_format` is unspecified
+    save_format = kwargs.get("save_format", "tf")
 
     # save keras save_format to path/data/save_format.txt
     with open(os.path.join(data_path, _KERAS_SAVE_FORMAT_PATH), "w") as f:


Test patch:
diff --git a/tests/keras/test_keras_model_export.py b/tests/keras/test_keras_model_export.py
--- a/tests/keras/test_keras_model_export.py
+++ b/tests/keras/test_keras_model_export.py
@@ -204,21 +204,23 @@ def _import_module(name, **kwargs):
         path0 = os.path.join(model_path, "0")
         with pytest.raises(MlflowException):
             mlflow.keras.save_model(x, path0)
-        mlflow.keras.save_model(x, path0, keras_module=FakeKerasModule)
+        mlflow.keras.save_model(x, path0, keras_module=FakeKerasModule, save_format="h5")
         y = mlflow.keras.load_model(path0)
         assert x == y
         path1 = os.path.join(model_path, "1")
-        mlflow.keras.save_model(x, path1, keras_module=FakeKerasModule.__name__)
+        mlflow.keras.save_model(x, path1, keras_module=FakeKerasModule.__name__, save_format="h5")
         z = mlflow.keras.load_model(path1)
         assert x == z
         # Tests model log
         with mlflow.start_run() as active_run:
             with pytest.raises(MlflowException):
                 mlflow.keras.log_model(x, "model0")
-            mlflow.keras.log_model(x, "model0", keras_module=FakeKerasModule)
+            mlflow.keras.log_model(x, "model0", keras_module=FakeKerasModule, save_format="h5")
             a = mlflow.keras.load_model("runs:/{}/model0".format(active_run.info.run_id))
             assert x == a
-            mlflow.keras.log_model(x, "model1", keras_module=FakeKerasModule.__name__)
+            mlflow.keras.log_model(
+                x, "model1", keras_module=FakeKerasModule.__name__, save_format="h5"
+            )
             b = mlflow.keras.load_model("runs:/{}/model1".format(active_run.info.run_id))
             assert x == b
 
@@ -495,13 +497,13 @@ def test_model_log_without_specified_conda_env_uses_default_env_with_expected_de
 
 @pytest.mark.large
 def test_model_load_succeeds_with_missing_data_key_when_data_exists_at_default_path(
-    model, model_path, data, predicted
+    tf_keras_model, model_path, data
 ):
     """
     This is a backwards compatibility test to ensure that models saved in MLflow version <= 0.8.0
     can be loaded successfully. These models are missing the `data` flavor configuration key.
     """
-    mlflow.keras.save_model(keras_model=model, path=model_path)
+    mlflow.keras.save_model(keras_model=tf_keras_model, path=model_path, save_format="h5")
     shutil.move(os.path.join(model_path, "data", "model.h5"), os.path.join(model_path, "model.h5"))
     model_conf_path = os.path.join(model_path, "MLmodel")
     model_conf = Model.load(model_conf_path)
@@ -511,7 +513,7 @@ def test_model_load_succeeds_with_missing_data_key_when_data_exists_at_default_p
     model_conf.save(model_conf_path)
 
     model_loaded = mlflow.keras.load_model(model_path)
-    assert all(model_loaded.predict(data[0].values) == predicted)
+    assert all(model_loaded.predict(data[0].values) == tf_keras_model.predict(data[0].values))
 
 
 @pytest.mark.release
diff --git a/tests/tensorflow_autolog/test_tensorflow2_autolog.py b/tests/tensorflow_autolog/test_tensorflow2_autolog.py
--- a/tests/tensorflow_autolog/test_tensorflow2_autolog.py
+++ b/tests/tensorflow_autolog/test_tensorflow2_autolog.py
@@ -2,6 +2,7 @@
 
 import collections
 import pytest
+from distutils.version import LooseVersion
 
 import numpy as np
 import pandas as pd
@@ -685,3 +686,63 @@ def test_flush_queue_is_thread_safe():
     flush_thread2.start()
     flush_thread2.join()
     assert len(mlflow.tensorflow._metric_queue) == 0
+
+
+def get_text_vec_model(train_samples):
+    # Taken from: https://github.com/mlflow/mlflow/issues/3910
+
+    from tensorflow.keras.layers.experimental.preprocessing import TextVectorization
+
+    VOCAB_SIZE = 10
+    SEQUENCE_LENGTH = 16
+    EMBEDDING_DIM = 16
+
+    vectorizer_layer = TextVectorization(
+        input_shape=(1,),
+        max_tokens=VOCAB_SIZE,
+        output_mode="int",
+        output_sequence_length=SEQUENCE_LENGTH,
+    )
+    vectorizer_layer.adapt(train_samples)
+    model = tf.keras.Sequential(
+        [
+            vectorizer_layer,
+            tf.keras.layers.Embedding(
+                VOCAB_SIZE, EMBEDDING_DIM, name="embedding", mask_zero=True, input_shape=(1,),
+            ),
+            tf.keras.layers.GlobalAveragePooling1D(),
+            tf.keras.layers.Dense(16, activation="relu"),
+            tf.keras.layers.Dense(1, activation="tanh"),
+        ]
+    )
+    model.compile(optimizer="adam", loss="mse", metrics="mae")
+    return model
+
+
+@pytest.mark.skipif(
+    LooseVersion(tf.__version__) < LooseVersion("2.3.0"),
+    reason=(
+        "Deserializing a model with `TextVectorization` and `Embedding`"
+        "fails in tensorflow < 2.3.0. See this issue:"
+        "https://github.com/tensorflow/tensorflow/issues/38250"
+    ),
+)
+def test_autolog_text_vec_model(tmpdir):
+    """
+    Verifies autolog successfully saves a model that can't be saved in the H5 format
+    """
+    mlflow.tensorflow.autolog()
+
+    train_samples = np.array(["this is an example", "another example"])
+    train_labels = np.array([0.4, 0.2])
+    model = get_text_vec_model(train_samples)
+
+    # Saving in the H5 format should fail
+    with pytest.raises(NotImplementedError, match="is not supported in h5"):
+        model.save(tmpdir.join("model.h5").strpath, save_format="h5")
+
+    with mlflow.start_run() as run:
+        model.fit(train_samples, train_labels, epochs=1)
+
+    loaded_model = mlflow.keras.load_model("runs:/" + run.info.run_id + "/model")
+    np.testing.assert_array_equal(loaded_model.predict(train_samples), model.predict(train_samples))