mlflow.keras

The mlflow.keras module provides an API for logging and loading Keras models. This module exports Keras models with the following flavors:

Keras (native) format
This is the main flavor that can be loaded back into Keras.
mlflow.pyfunc
Produced for use by generic pyfunc-based deployment tools and batch inference.
mlflow.keras.load_model(path, run_id=None)

Load a Keras model from a local file (if run_id is None) or a run.

Parameters:
  • path – Local filesystem path or run-relative artifact path to the model saved by mlflow.keras.log_model().
  • run_id – Run ID. If provided, combined with path to identify the model.
 
>>> # Load persisted model as a Keras model or as a PyFunc, call predict() on a Pandas DataFrame
>>> keras_model = mlflow.keras.load_model("models", run_id="96771d893a5e46159d9f3b49bf9013e2")
>>> predictions = keras_model.predict(x_test)
mlflow.keras.log_model(keras_model, artifact_path, **kwargs)

Log a Keras model as an MLflow artifact for the current run.

Parameters:
  • keras_model – Keras model to be saved.
  • artifact_path – Run-relative artifact path.
  • kwargs – kwargs to pass to keras_model.save method.
 
>>> from keras import Dense, layers
>>> import mlflow
>>> # Build, compile, and train your model
>>> keras_model = ...
>>> keras_model.compile(optimizer="rmsprop", loss="mse", metrics["accuracy"])
>>> results = keras_model.fit(
...     x_train, y_train, epochs=20, batch_size = 128, validation_data=(x_val, y_val))
>>> # Log metrics and log the model
>>> with mlflow.start_run() as run:
>>>   mlflow.keras.log_model(keras_model, "models")
mlflow.keras.save_model(keras_model, path, conda_env=None, mlflow_model=<mlflow.models.Model object>)

Save a Keras model to a path on the local file system.

Parameters:
  • keras_model – Keras model to be saved.
  • path – Local path where the model is to be saved.
  • mlflow_model – MLflow model config this flavor is being added to.
 
>>> import mlflow
>>> # Build, compile, and train your model
>>> keras_model = ...
>>> keras_model_path = ...
>>> keras_model.compile(optimizer="rmsprop", loss="mse", metrics["accuracy"])
>>> results = keras_model.fit(
...     x_train, y_train, epochs=20, batch_size = 128, validation_data=(x_val, y_val))
... # Save the model as an MLflow Model
>>> mlflow.keras.save_model(keras_model, keras_model_path)