Repository: mlflow/mlflow

Issue:
[BUG] Cannot save Keras/TF_2.x model as SavedModel format using mlflow.keras.log_model with kwargs { "save_format": "tf" }
Thank you for submitting an issue. Please refer to our [issue policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md) for additional information about bug reports. For help with debugging your code, please refer to [Stack Overflow](https://stackoverflow.com/questions/tagged/mlflow).

**Please fill in this bug report template to ensure a timely and thorough response.**

### 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.10.0
Keras version: 2.4.0
TensorFlow version: 2.3.0
Python Version: 3.7.6
Operating System: Darwin - 19.4.0

### Describe the problem

Cannot save TensorFlow 2.3.0  Keras model as TensorFlow SaveModel format  using mlflow.keras.log_model since it does not honor kwargs { "save_format": "tf" }.  The saved model is still H5.  The  suffix is [hard-coded to h5](https://github.com/mlflow/mlflow/blob/master/mlflow/keras.py#L207) and not dependent on the value of save_format.

TensorFlow SaveModel format is now the default for Keras TensorFlow 2.

### Code to reproduce issue
```
import mlflow
import mlflow.keras
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

def baseline_model():
    model = Sequential()
    model.add(Dense(2, input_dim=2, kernel_initializer="normal", activation="relu"))
    model.compile(loss="mean_squared_error", optimizer="adam")
    return model
model = baseline_model()
with mlflow.start_run() as run:
    mlflow.keras.log_model(model, "tensorflow-model", **{ "save_format": "tf" })
```

### Other info / logs

`{ "save_format": "tf" }`  is correctly passed to `keras_model.save`  at https://github.com/mlflow/mlflow/blob/master/mlflow/keras.py#L208 but output is H5.

```
        with tempfile.NamedTemporaryFile(suffix='.h5') as f:
            keras_model.save(f.name, **kwargs)
            f.flush()  # force flush the data
            shutil.copyfile(src=f.name, dst=model_path)
```

### What component(s), interfaces, languages, and integrations does this bug affect?
Components 
- [ x] `area/models`: MLmodel format, model serialization/deserialization, flavors

###  Related Issues
- https://github.com/mlflow/mlflow/issues/3246 - [FR] Serialization format for TensorFlow 2.x (Keras) model should be Save Model per TF 2.0 doc



Discussion:
@amesar Thanks for filing this. Would you be interested in making a PR?
Fixing this issue might break backward compatibility because `mlflow.kears` currently expects a keras model to be saved as `model.h5`, but to save a model in the `SavedModel` format, we have to use a file name that **does not end with ".h5"** (e.g. `model`). This means older versions of MLflow would fail to load models logged by newer versions of MLflow in the `SaveModel` format.
Hi, is there a workaround for this issue?
I think it's worth implementing, as going towards the `SavedModel` format opens up a lot of possibilities in the TensorFlow space, particularly TFX and TF Serving.  I don't think there is a solution that maintains backwards compatibility, but its worth a shot.  Is anyone working on this currently?
@harupy I'm happy to contribute a fix here; but would like to discuss how to move forward on breaking backwards compat. One option that might be easiest on users (but maybe leaves a bit to clean up later):

1. Change saving to not use the h5 prefix
2. When loading, check if file exists; if not fallback to looking for file with h5 prefix (to maintain backwards compat)

This way, backwards compat can be maintained (though with a slight perf hit for model loading -- I don't believe it's a critical perf difference though).

Then later when you want to make it a breaking change (e.g. next time you do a significant version update), we can remove the fallback.

Thoughts?

EDIT: It looks like someone has taken a stab at [fixing this bug](https://github.com/mlflow/mlflow/pull/3282); I think it still needs some work, but looks like the train is moving on this one.

Resolving patch:
diff --git a/mlflow/keras.py b/mlflow/keras.py
--- a/mlflow/keras.py
+++ b/mlflow/keras.py
@@ -36,8 +36,9 @@
 # File name to which custom objects cloudpickle is saved - used during save and load
 _CUSTOM_OBJECTS_SAVE_PATH = "custom_objects.cloudpickle"
 _KERAS_MODULE_SPEC_PATH = "keras_module.txt"
+_KERAS_SAVE_FORMAT_PATH = "save_format.txt"
 # File name to which keras model is saved
-_MODEL_SAVE_PATH = "model.h5"
+_MODEL_SAVE_PATH = "model"
 # Conda env subpath when saving/loading model
 _CONDA_ENV_SUBPATH = "conda.yaml"
 
@@ -90,7 +91,7 @@ def save_model(
     keras_module=None,
     signature: ModelSignature = None,
     input_example: ModelInputExample = None,
-    **kwargs
+    **kwargs,
 ):
     """
     Save a Keras model to a path on the local file system.
@@ -215,9 +216,21 @@ def _is_tf_keras(model):
     with open(os.path.join(data_path, _KERAS_MODULE_SPEC_PATH), "w") as f:
         f.write(keras_module.__name__)
 
-    # save keras model to path/data/model.h5
+    # 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")
+
+    # save keras save_format to path/data/save_format.txt
+    with open(os.path.join(data_path, _KERAS_SAVE_FORMAT_PATH), "w") as f:
+        f.write(save_format)
+
+    # save keras model
+    # To maintain prior behavior, when the format is HDF5, we save
+    # with the h5 file extension. Otherwise, model_path is a directory
+    # where the saved_model.pb will be stored (for SavedModel format)
+    file_extension = ".h5" if save_format == "h5" else ""
     model_subpath = os.path.join(data_subpath, _MODEL_SAVE_PATH)
-    model_path = os.path.join(path, model_subpath)
+    model_path = os.path.join(path, model_subpath) + file_extension
     if path.startswith("/dbfs/"):
         # The Databricks Filesystem uses a FUSE implementation that does not support
         # random writes. It causes an error.
@@ -233,6 +246,7 @@ def _is_tf_keras(model):
         FLAVOR_NAME,
         keras_module=keras_module.__name__,
         keras_version=keras_module.__version__,
+        save_format=save_format,
         data=data_subpath,
     )
 
@@ -266,7 +280,7 @@ def log_model(
     signature: ModelSignature = None,
     input_example: ModelInputExample = None,
     await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
-    **kwargs
+    **kwargs,
 ):
     """
     Log a Keras model as an MLflow artifact for the current run.
@@ -352,7 +366,7 @@ def log_model(
         signature=signature,
         input_example=input_example,
         await_registration_for=await_registration_for,
-        **kwargs
+        **kwargs,
     )
 
 
@@ -375,7 +389,7 @@ def _save_custom_objects(path, custom_objects):
         cloudpickle.dump(custom_objects, out_f)
 
 
-def _load_model(model_path, keras_module, **kwargs):
+def _load_model(model_path, keras_module, save_format, **kwargs):
     keras_models = importlib.import_module(keras_module.__name__ + ".models")
     custom_objects = kwargs.pop("custom_objects", {})
     custom_objects_path = None
@@ -390,9 +404,17 @@ def _load_model(model_path, keras_module, **kwargs):
             pickled_custom_objects = cloudpickle.load(in_f)
             pickled_custom_objects.update(custom_objects)
             custom_objects = pickled_custom_objects
+
+    # If the save_format is HDF5, then we save with h5 file
+    # extension to align with prior behavior of mlflow logging
+    if save_format == "h5":
+        model_path = model_path + ".h5"
+
     from distutils.version import StrictVersion
 
-    if StrictVersion(keras_module.__version__.split("-")[0]) >= StrictVersion("2.2.3"):
+    if save_format == "h5" and StrictVersion(
+        keras_module.__version__.split("-")[0]
+    ) >= StrictVersion("2.2.3"):
         # NOTE: Keras 2.2.3 does not work with unicode paths in python2. Pass in h5py.File instead
         # of string to avoid issues.
         import h5py
@@ -439,6 +461,15 @@ def _load_pyfunc(path):
 
         keras_module = keras
 
+    # By default, we assume the save_format is h5 for backwards compatibility
+    save_format = "h5"
+    save_format_path = os.path.join(path, _KERAS_SAVE_FORMAT_PATH)
+    if os.path.isfile(save_format_path):
+        with open(save_format_path, "r") as f:
+            save_format = f.read()
+
+    # In SavedModel format, if we don't compile the model
+    should_compile = save_format == "tf"
     K = importlib.import_module(keras_module.__name__ + ".backend")
     if keras_module.__name__ == "tensorflow.keras" or K.backend() == "tensorflow":
         if LooseVersion(tf.__version__) < LooseVersion("2.0.0"):
@@ -450,11 +481,18 @@ def _load_pyfunc(path):
             with graph.as_default():
                 with sess.as_default():  # pylint:disable=not-context-manager
                     K.set_learning_phase(0)
-                    m = _load_model(path, keras_module=keras_module, compile=False)
+                    m = _load_model(
+                        path,
+                        keras_module=keras_module,
+                        save_format=save_format,
+                        compile=should_compile,
+                    )
                     return _KerasModelWrapper(m, graph, sess)
         else:
             K.set_learning_phase(0)
-            m = _load_model(path, keras_module=keras_module, compile=False)
+            m = _load_model(
+                path, keras_module=keras_module, save_format=save_format, compile=should_compile
+            )
             return _KerasModelWrapper(m, None, None)
 
     else:
@@ -495,7 +533,14 @@ def load_model(model_uri, **kwargs):
     keras_model_artifacts_path = os.path.join(
         local_model_path, flavor_conf.get("data", _MODEL_SAVE_PATH)
     )
-    return _load_model(model_path=keras_model_artifacts_path, keras_module=keras_module, **kwargs)
+    # For backwards compatibility, we assume h5 when the save_format is absent
+    save_format = flavor_conf.get("save_format", "h5")
+    return _load_model(
+        model_path=keras_model_artifacts_path,
+        keras_module=keras_module,
+        save_format=save_format,
+        **kwargs,
+    )
 
 
 @experimental


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
@@ -210,10 +210,14 @@ def _import_module(name, **kwargs):
             assert x == b
 
 
-@pytest.mark.skip("Unskip this test once https://github.com/h5py/h5py/issues/1732 is fixed")
-@pytest.mark.parametrize("build_model", [model, tf_keras_model])
+@pytest.mark.parametrize(
+    "build_model,save_format",
+    # TODO: Unskip these parameters once https://github.com/h5py/h5py/issues/1732 is fixed
+    # [(model, None), (tf_keras_model, None), (tf_keras_model, "h5"),
+    [(tf_keras_model, "tf")],
+)
 @pytest.mark.large
-def test_model_save_load(build_model, model_path, data):
+def test_model_save_load(build_model, save_format, model_path, data):
     x, _ = data
     keras_model = build_model(data)
     if build_model == tf_keras_model:
@@ -221,14 +225,19 @@ def test_model_save_load(build_model, model_path, data):
     else:
         model_path = os.path.join(model_path, "plain")
     expected = keras_model.predict(x)
-    mlflow.keras.save_model(keras_model, model_path)
+    kwargs = {"save_format": save_format} if save_format else {}
+    mlflow.keras.save_model(keras_model, model_path, **kwargs)
     # Loading Keras model
     model_loaded = mlflow.keras.load_model(model_path)
-    assert type(keras_model) == type(model_loaded)
-    assert all(expected == model_loaded.predict(x))
+    # When saving as SavedModel, we actually convert the model
+    # to a slightly different format, so we cannot assume it is
+    # exactly the same.
+    if save_format != "tf":
+        assert type(keras_model) == type(model_loaded)
+    np.testing.assert_allclose(model_loaded.predict(x), expected, rtol=1e-5)
     # Loading pyfunc model
     pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
-    assert all(pyfunc_loaded.predict(x).values == expected)
+    np.testing.assert_allclose(pyfunc_loaded.predict(x).values, expected, rtol=1e-5)
 
     # pyfunc serve
     scoring_response = pyfunc_serve_and_score_model(
@@ -237,12 +246,11 @@ def test_model_save_load(build_model, model_path, data):
         content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON_SPLIT_ORIENTED,
     )
     print(scoring_response.content)
-    assert all(
-        pd.read_json(scoring_response.content, orient="records", encoding="utf8").values.astype(
-            np.float32
-        )
-        == expected
-    )
+    actual_scoring_response = pd.read_json(
+        scoring_response.content, orient="records", encoding="utf8"
+    ).values.astype(np.float32)
+    np.testing.assert_allclose(actual_scoring_response, expected, rtol=1e-5)
+
     # test spark udf
     spark_udf_preds = score_model_as_udf(
         model_uri=os.path.abspath(model_path), pandas_df=pd.DataFrame(x), result_type="float"
@@ -508,3 +516,53 @@ def test_sagemaker_docker_model_scoring_with_default_conda_env(model, model_path
     deployed_model_preds = pd.DataFrame(json.loads(scoring_response.content))
 
     np.testing.assert_array_almost_equal(deployed_model_preds.values, predicted, decimal=4)
+
+
+def test_save_model_with_tf_save_format(model_path):
+    """Ensures that Keras models can be saved with SavedModel format.
+
+    Using SavedModel format (save_format="tf") requires that the file extension
+    is _not_ "h5".
+    """
+    keras_model = mock.Mock(spec=tf.keras.Model)
+    mlflow.keras.save_model(keras_model=keras_model, path=model_path, save_format="tf")
+    _, args, kwargs = keras_model.save.mock_calls[0]
+    # Ensure that save_format propagated through
+    assert kwargs["save_format"] == "tf"
+    # Ensure that the saved model does not have h5 extension
+    assert not args[0].endswith(".h5")
+
+
+@pytest.mark.large
+def test_save_and_load_model_with_tf_save_format(tf_keras_model, model_path):
+    """Ensures that keras models saved with save_format="tf" can be loaded."""
+    mlflow.keras.save_model(keras_model=tf_keras_model, path=model_path, save_format="tf")
+    model_conf_path = os.path.join(model_path, "MLmodel")
+    model_conf = Model.load(model_conf_path)
+    flavor_conf = model_conf.flavors.get(mlflow.keras.FLAVOR_NAME, None)
+    assert flavor_conf is not None
+    assert flavor_conf.get("save_format") == "tf"
+    assert not os.path.exists(
+        os.path.join(model_path, "data", "model.h5")
+    ), "TF model was saved with HDF5 format; expected SavedModel"
+    assert os.path.isdir(
+        os.path.join(model_path, "data", "model")
+    ), "Expected directory containing saved_model.pb"
+
+    model_loaded = mlflow.keras.load_model(model_path)
+    assert tf_keras_model.to_json() == model_loaded.to_json()
+
+
+@pytest.mark.large
+def test_load_without_save_format(tf_keras_model, model_path):
+    """Ensures that keras models without save_format can still be loaded."""
+    mlflow.keras.save_model(tf_keras_model, model_path, save_format="h5")
+    model_conf_path = os.path.join(model_path, "MLmodel")
+    model_conf = Model.load(model_conf_path)
+    flavor_conf = model_conf.flavors.get(mlflow.keras.FLAVOR_NAME)
+    assert flavor_conf is not None
+    del flavor_conf["save_format"]
+    model_conf.save(model_conf_path)
+
+    model_loaded = mlflow.keras.load_model(model_path)
+    assert tf_keras_model.to_json() == model_loaded.to_json()