Repository: mlflow/mlflow

Issue:
Support for 'compile=False' when loading Keras models
### System information
- **Have I written custom code (as opposed to using a stock example script provided in MLflow)**: Yes
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:
- **MLflow installed from (source or binary)**:  binary
- **MLflow version (run ``mlflow --version``)**: 0.9
- **Python version**: 3.6
- **npm version (if running the dev UI):
- **Exact command to reproduce**:

### Describe the problem

When loading Keras models that were trained with custom losses and metrics, Keras will throw an error if those functions are not found:
https://github.com/keras-team/keras/issues/5916

These functions are not needed for prediction, and this can be solved by adding the 'compile=False' keyword parameter when calling load_model:
https://github.com/keras-team/keras/issues/5916#issuecomment-457624404

However, mlflow does not support passing any of these when loading and saving Keras models. 

Support for this should be added to support more complex models and training use cases.

In the meanwhile, does anyone have any suggestions for workarounds?


Support for 'compile=False' when loading Keras models
### System information
- **Have I written custom code (as opposed to using a stock example script provided in MLflow)**: Yes
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:
- **MLflow installed from (source or binary)**:  binary
- **MLflow version (run ``mlflow --version``)**: 0.9
- **Python version**: 3.6
- **npm version (if running the dev UI):
- **Exact command to reproduce**:

### Describe the problem

When loading Keras models that were trained with custom losses and metrics, Keras will throw an error if those functions are not found:
https://github.com/keras-team/keras/issues/5916

These functions are not needed for prediction, and this can be solved by adding the 'compile=False' keyword parameter when calling load_model:
https://github.com/keras-team/keras/issues/5916#issuecomment-457624404

However, mlflow does not support passing any of these when loading and saving Keras models. 

Support for this should be added to support more complex models and training use cases.

In the meanwhile, does anyone have any suggestions for workarounds?




Discussion:
The issue can be resolved by changing this function in keras.py:

```
def _load_model(model_file):
    import keras.models
    import h5py

    from distutils.version import StrictVersion

    if StrictVersion(keras.__version__) >= 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.
        with h5py.File(os.path.abspath(model_file), "r") as model_file:
            return keras.models.load_model(model_file, compile=False)
    else:
        # NOTE: Older versions of Keras only handle filepath.
        return keras.models.load_model(model_file)
```

The additional flags supported by Keras should ideally be passed in to save_model() and _load_model() via `**kwargs`.


The issue can be resolved by changing this function in keras.py:

```
def _load_model(model_file):
    import keras.models
    import h5py

    from distutils.version import StrictVersion

    if StrictVersion(keras.__version__) >= 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.
        with h5py.File(os.path.abspath(model_file), "r") as model_file:
            return keras.models.load_model(model_file, compile=False)
    else:
        # NOTE: Older versions of Keras only handle filepath.
        return keras.models.load_model(model_file)
```

The additional flags supported by Keras should ideally be passed in to save_model() and _load_model() via `**kwargs`.



Resolving patch:
diff --git a/mlflow/keras.py b/mlflow/keras.py
--- a/mlflow/keras.py
+++ b/mlflow/keras.py
@@ -146,21 +146,21 @@ def log_model(keras_model, artifact_path, conda_env=None, **kwargs):
               keras_model=keras_model, conda_env=conda_env, **kwargs)
 
 
-def _load_model(model_file):
+def _load_model(model_file, **kwargs):
     import keras
     import keras.models
     import h5py
 
     from distutils.version import StrictVersion
 
-    if StrictVersion(keras.__version__) >= StrictVersion("2.2.3"):
+    if StrictVersion(keras.__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.
         with h5py.File(os.path.abspath(model_file), "r") as model_file:
-            return keras.models.load_model(model_file)
+            return keras.models.load_model(model_file, **kwargs)
     else:
         # NOTE: Older versions of Keras only handle filepath.
-        return keras.models.load_model(model_file)
+        return keras.models.load_model(model_file, **kwargs)
 
 
 class _KerasModelWrapper:
@@ -195,15 +195,16 @@ 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)
+                m = _load_model(path, compile=False)
         return _KerasModelWrapper(m, graph, sess)
     else:
         raise Exception("Unsupported backend '%s'" % K._BACKEND)
 
 
-def load_model(model_uri):
+def load_model(model_uri, **kwargs):
     """
     Load a Keras model from a local file (if ``run_id`` is None) or a run.
+    Extra arguments are passed through to keras.load_model.
 
     :param model_uri: The location, in URI format, of the MLflow model, for example:
 
@@ -227,4 +228,4 @@ def load_model(model_uri):
     # Flavor configurations for models saved in MLflow version <= 0.8.0 may not contain a
     # `data` key; in this case, we assume the model artifact path to be `model.h5`
     keras_model_artifacts_path = os.path.join(local_model_path, flavor_conf.get("data", "model.h5"))
-    return _load_model(model_file=keras_model_artifacts_path)
+    return _load_model(model_file=keras_model_artifacts_path, **kwargs)


Test patch: