Repository: mlflow/mlflow

Issue:
mlflow project can't pickle
### System information
- **Have I written custom code (as opposed to using a stock example script provided in MLflow)**:
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:
- **MLflow installed from (source or binary)**: 
- **MLflow version (run ``mlflow --version``)**:
- **Python version**: 
- **npm version (if running the dev UI):
- **Exact command to reproduce**:

### Describe the problem
when i run my mfflow project ,the python  raise TypeError:"can't pickle %s objects" % base.__name__
TypeError: can't pickle module objects
I don't know how to solve it

### Source code / logs
Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks,
please include the full traceback. Large logs and files should be attached.
Try to provide a reproducible test case that is the bare minimum necessary to generate the problem.



Discussion:
Hi @panxuexue, thanks for reporting this issue - could you post the code you're trying to run & the full Python stacktrace? Thanks!
python code
import numpy as np

from keras.models import Sequential
from keras.layers import Dense
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import mlflow.sklearn
import sys
reload(sys)
sys.setdefaultencoding('utf8')
np.random.seed(1337)  # for reproducibility

if __name__ == "__main__":

    X = np.linspace(-1, 1, 200)
    np.random.shuffle(X)  # randomize the data
    Y = 0.5 * X + 2 + np.random.normal(0, 0.05, (200,))
    # plot data
    # plt.scatter(X, Y)
    # plt.show()
    X_train, Y_train = X[:160], Y[:160]  # first 160 data points
    X_test, Y_test = X[160:], Y[160:]  # last 40 data points

    # build a neural network from the 1st layer to the last layer
    with mlflow.start_run():
        # build a neural network from the 1st layer to the last layer
        model = Sequential()
        model.add(Dense(output_dim=1, input_dim=1))

        # choose loss function and optimizing method
        loss = str(sys.argv[1]) if len(sys.argv) > 1 else 'mse'
        optimizer = str(sys.argv[2]) if len(sys.argv) > 2 else 'sgd'

        model.compile(loss=loss, optimizer=optimizer)

        # training
        print('Training -----------')
        for step in range(301):
            cost = model.train_on_batch(X_train, Y_train)
            if step % 100 == 0:
                print('train cost: ', cost)

        # test
        print('\nTesting ------------')
        cost = model.evaluate(X_test, Y_test, batch_size=40)
        print('test cost:', cost)
        W, b = model.layers[0].get_weights()
        print('Weights=', W, '\nbiases=', b)
        Y_pred = model.predict(X_test)

        rmse = np.sqrt(mean_squared_error(Y_test, Y_pred))

        mae = mean_absolute_error(Y_test, Y_pred)

        r2 = r2_score(Y_test, Y_pred)

        # plotting the prediction

        # plt.scatter(X_test, Y_test)
        #
        # plt.plot(X_test, Y_pred)

        # plt.show()

        print("Sequential model (loss=%s, optimizer=%s):" % (loss, optimizer))
        mlflow.log_param("loss", loss)
        mlflow.log_param("optimizer", optimizer)
        mlflow.log_metric("rmse", rmse)
        mlflow.log_metric("r2", r2)
        mlflow.log_metric("mae", mae)
        mlflow.log_metric("weight", W)
        mlflow.log_metric("biases", b)
        mlflow.sklearn.log_model(model, "model")


the MLproject
name: text_mlflow

entry_points:
  main:
    parameters:
      loss: {type: string,default: 'mse'}
      optimizer: {type: string, default: 'sgd'}
    command: "python reg.py {loss} {optimizer}"


I don't use conda


Thanks!


Resolving patch:
diff --git a/mlflow/keras.py b/mlflow/keras.py
new file mode 100644
--- /dev/null
+++ b/mlflow/keras.py
@@ -0,0 +1,68 @@
+"""Sample MLflow integration for Keras."""
+
+from __future__ import absolute_import
+
+import os
+
+from mlflow import pyfunc
+from mlflow.models import Model
+import mlflow.tracking
+
+import pandas as pd
+
+
+def save_model(keras_model, path, conda_env=None, mlflow_model=Model()):
+    """
+    Save a Keras model to a path on the local file system.
+
+    :param keras_model: Keras model to be saved.
+    :param path: Local path where the model is to be saved.
+    :param mlflow_model: MLflow model config this flavor is being added to.
+    """
+    import keras
+
+    path = os.path.abspath(path)
+    if os.path.exists(path):
+        raise Exception("Path '{}' already exists".format(path))
+    os.makedirs(path)
+    model_file = os.path.join(path, "model.h5")
+    keras_model.save(model_file)
+
+    pyfunc.add_to_model(mlflow_model, loader_module="mlflow.keras",
+                        data="model.h5", env=conda_env)
+    mlflow_model.add_flavor("keras", keras_version=keras.__version__)
+    mlflow_model.save(os.path.join(path, "MLmodel"))
+
+
+def log_model(keras_model, artifact_path, **kwargs):
+    """Log a Keras model as an MLflow artifact for the current run."""
+    Model.log(artifact_path=artifact_path, flavor=mlflow.keras,
+              keras_model=keras_model, **kwargs)
+
+
+def _load_model(model_file):
+    import keras.models
+    return keras.models.load_model(os.path.abspath(model_file))
+
+
+class _KerasModelWrapper:
+    def __init__(self, keras_model):
+        self.keras_model = keras_model
+
+    def predict(self, dataframe):
+        predicted = pd.DataFrame(self.keras_model.predict(dataframe))
+        predicted.index = dataframe.index
+        return predicted
+
+
+def load_pyfunc(model_file):
+    return _KerasModelWrapper(_load_model(model_file))
+
+
+def load_model(path, run_id=None):
+    """
+    Load a Keras model from a local file (if run_id is None) or a run.
+    """
+    if run_id is not None:
+        path = mlflow.tracking._get_model_log_dir(model_name=path, run_id=run_id)
+    return _load_model(os.path.join(path, "model.h5"))


Test patch:
diff --git a/tests/keras/test_keras_model_export.py b/tests/keras/test_keras_model_export.py
new file mode 100644
--- /dev/null
+++ b/tests/keras/test_keras_model_export.py
@@ -0,0 +1,85 @@
+# pep8: disable=E501
+
+from __future__ import print_function
+
+import os
+import pytest
+from keras.models import Sequential
+from keras.layers import Dense
+import sklearn.datasets as datasets
+import pandas as pd
+import numpy as np
+
+import mlflow.keras
+from mlflow import pyfunc
+from mlflow import tracking
+from mlflow.utils.file_utils import TempDir
+
+
+@pytest.fixture(scope='module')
+def data():
+    iris = datasets.load_iris()
+    data = pd.DataFrame(data=np.c_[iris['data'], iris['target']],
+                        columns=iris['feature_names'] + ['target'])
+    y = data['target']
+    x = data.drop('target', axis=1)
+    return (x, y)
+
+
+@pytest.fixture(scope='module')
+def model(data):
+    x, y = data
+    model = Sequential()
+    model.add(Dense(3, input_dim=4))
+    model.add(Dense(1))
+    model.compile(loss='mean_squared_error', optimizer='SGD')
+    model.fit(x, y)
+    return model
+
+
+@pytest.fixture(scope='module')
+def predicted(model, data):
+    return model.predict(data[0])
+
+
+def test_model_save_load(model, data, predicted):
+    x, y = data
+    with TempDir(chdr=True, remove_on_exit=True) as tmp:
+        path = tmp.path("model")
+        mlflow.keras.save_model(model, path)
+
+        # Loading Keras model
+        model_loaded = mlflow.keras.load_model(path)
+        assert all(model_loaded.predict(x) == predicted)
+
+        # Loading pyfunc model
+        pyfunc_loaded = mlflow.pyfunc.load_pyfunc(path)
+        assert all(pyfunc_loaded.predict(x).values == predicted)
+
+
+def test_model_log(model, data, predicted):
+    x, y = data
+    old_uri = tracking.get_tracking_uri()
+    # should_start_run tests whether or not calling log_model() automatically starts a run.
+    for should_start_run in [False, True]:
+        with TempDir(chdr=True, remove_on_exit=True) as tmp:
+            try:
+                tracking.set_tracking_uri("test")
+                if should_start_run:
+                    tracking.start_run()
+                mlflow.keras.log_model(model, artifact_path="keras_model")
+
+                # Load model
+                model_loaded = mlflow.keras.load_model(
+                    "keras_model",
+                    run_id=tracking.active_run().info.run_uuid)
+                assert all(model_loaded.predict(x) == predicted)
+
+                # Loading pyfunc model
+                pyfunc_loaded = mlflow.pyfunc.load_pyfunc(
+                    "keras_model",
+                    run_id=tracking.active_run().info.run_uuid)
+                assert all(pyfunc_loaded.predict(x).values == predicted)
+            finally:
+                tracking.end_run()
+    tracking.set_tracking_uri(old_uri)