MLflow integration for TensorFlow.

Manages logging and loading TensorFlow models as Python Functions. You are expected to save your own saved_models and pass their paths to log_saved_model() so that MLflow can track the models.

In order to load the model to predict on it again, you can call model = mlflow.pyfunc.load_pyfunc(saved_model_dir), followed by prediction = model.predict(pandas DataFrame) to obtain a prediction in a pandas DataFrame.

The loaded PyFunc model does not expose any APIs for model training.


Load model stored in Python Function format. The loaded model object exposes a predict(pandas DataFrame) method that returns a Pandas DataFrame containing the model’s inference output on an input DataFrame.

Parameters:saved_model_dir – Directory where the model is saved.
Return type:Pyfunc format model with function model.predict(pandas DataFrame) -> pandas DataFrame).
mlflow.tensorflow.log_saved_model(saved_model_dir, signature_def_key, artifact_path)

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

  • saved_model_dir – Directory where the exported TensorFlow model is saved.
  • signature_def_key – The signature definition to use when loading the model again. See SignatureDefs in SavedModel for TensorFlow Serving for details.
  • artifact_path – Path (within the artifact directory for the current run) to which artifacts of the model will be saved.