mlflow.sklearn

The mlflow.sklearn module provides an API for logging and loading scikit-learn models. This module exports scikit-learn models with the following flavors:

Python (native) pickle format

This is the main flavor that can be loaded back into scikit-learn.

mlflow.pyfunc

Produced for use by generic pyfunc-based deployment tools and batch inference. NOTE: The mlflow.pyfunc flavor is only added for scikit-learn models that define predict(), since predict() is required for pyfunc model inference.

mlflow.sklearn.autolog(log_input_examples=False, log_model_signatures=True, log_models=True, log_datasets=True, disable=False, exclusive=False, disable_for_unsupported_versions=False, silent=False, max_tuning_runs=5, log_post_training_metrics=True, serialization_format='cloudpickle', registered_model_name=None, pos_label=None, extra_tags=None)[source]

Note

Autologging is known to be compatible with the following package versions: 0.24.1 <= scikit-learn <= 1.5.2. Autologging may not succeed when used with package versions outside of this range.

Enables (or disables) and configures autologging for scikit-learn estimators.

When is autologging performed?

Autologging is performed when you call:

  • estimator.fit()

  • estimator.fit_predict()

  • estimator.fit_transform()

Logged information
Parameters

  • Parameters obtained by estimator.get_params(deep=True). Note that get_params is called with deep=True. This means when you fit a meta estimator that chains a series of estimators, the parameters of these child estimators are also logged.

Training metrics
Post training metrics

When users call metric APIs after model training, MLflow tries to capture the metric API results and log them as MLflow metrics to the Run associated with the model. The following types of scikit-learn metric APIs are supported:

  • model.score

  • metric APIs defined in the sklearn.metrics module

For post training metrics autologging, the metric key format is: “{metric_name}[-{call_index}]_{dataset_name}”

  • If the metric function is from sklearn.metrics, the MLflow “metric_name” is the metric function name. If the metric function is model.score, then “metric_name” is “{model_class_name}_score”.

  • If multiple calls are made to the same scikit-learn metric API, each subsequent call adds a “call_index” (starting from 2) to the metric key.

  • MLflow uses the prediction input dataset variable name as the “dataset_name” in the metric key. The “prediction input dataset variable” refers to the variable which was used as the first argument of the associated model.predict or model.score call. Note: MLflow captures the “prediction input dataset” instance in the outermost call frame and fetches the variable name in the outermost call frame. If the “prediction input dataset” instance is an intermediate expression without a defined variable name, the dataset name is set to “unknown_dataset”. If multiple “prediction input dataset” instances have the same variable name, then subsequent ones will append an index (starting from 2) to the inspected dataset name.

Limitations

  • MLflow can only map the original prediction result object returned by a model prediction API (including predict / predict_proba / predict_log_proba / transform, but excluding fit_predict / fit_transform.) to an MLflow run. MLflow cannot find run information for other objects derived from a given prediction result (e.g. by copying or selecting a subset of the prediction result). scikit-learn metric APIs invoked on derived objects do not log metrics to MLflow.

  • Autologging must be enabled before scikit-learn metric APIs are imported from sklearn.metrics. Metric APIs imported before autologging is enabled do not log metrics to MLflow runs.

  • If user define a scorer which is not based on metric APIs in sklearn.metrics, then then post training metric autologging for the scorer is invalid.

Tags

  • An estimator class name (e.g. “LinearRegression”).

  • A fully qualified estimator class name (e.g. “sklearn.linear_model._base.LinearRegression”).

Artifacts

  • An MLflow Model with the mlflow.sklearn flavor containing a fitted estimator (logged by mlflow.sklearn.log_model()). The Model also contains the mlflow.pyfunc flavor when the scikit-learn estimator defines predict().

  • For post training metrics API calls, a “metric_info.json” artifact is logged. This is a JSON object whose keys are MLflow post training metric names (see “Post training metrics” section for the key format) and whose values are the corresponding metric call commands that produced the metrics, e.g. accuracy_score(y_true=test_iris_y, y_pred=pred_iris_y, normalize=False).

How does autologging work for meta estimators?

When a meta estimator (e.g. Pipeline, GridSearchCV) calls fit(), it internally calls fit() on its child estimators. Autologging does NOT perform logging on these constituent fit() calls.

Parameter search

In addition to recording the information discussed above, autologging for parameter search meta estimators (GridSearchCV and RandomizedSearchCV) records child runs with metrics for each set of explored parameters, as well as artifacts and parameters for the best model (if available).

Supported estimators

Example

See more examples

from pprint import pprint
import numpy as np
from sklearn.linear_model import LinearRegression
import mlflow
from mlflow import MlflowClient


def fetch_logged_data(run_id):
    client = MlflowClient()
    data = client.get_run(run_id).data
    tags = {k: v for k, v in data.tags.items() if not k.startswith("mlflow.")}
    artifacts = [f.path for f in client.list_artifacts(run_id, "model")]
    return data.params, data.metrics, tags, artifacts


# enable autologging
mlflow.sklearn.autolog()

# prepare training data
X = np.array([[1, 1], [1, 2], [2, 2], [2, 3]])
y = np.dot(X, np.array([1, 2])) + 3

# train a model
model = LinearRegression()
with mlflow.start_run() as run:
    model.fit(X, y)

# fetch logged data
params, metrics, tags, artifacts = fetch_logged_data(run.info.run_id)

pprint(params)
# {'copy_X': 'True',
#  'fit_intercept': 'True',
#  'n_jobs': 'None',
#  'normalize': 'False'}

pprint(metrics)
# {'training_score': 1.0,
#  'training_mean_absolute_error': 2.220446049250313e-16,
#  'training_mean_squared_error': 1.9721522630525295e-31,
#  'training_r2_score': 1.0,
#  'training_root_mean_squared_error': 4.440892098500626e-16}

pprint(tags)
# {'estimator_class': 'sklearn.linear_model._base.LinearRegression',
#  'estimator_name': 'LinearRegression'}

pprint(artifacts)
# ['model/MLmodel', 'model/conda.yaml', 'model/model.pkl']
Parameters

  • log_input_examples – If True, input examples from training datasets are collected and logged along with scikit-learn model artifacts during training. If False, input examples are not logged. Note: Input examples are MLflow model attributes and are only collected if log_models is also True.

  • log_model_signatures – If True, ModelSignatures describing model inputs and outputs are collected and logged along with scikit-learn model artifacts during training. If False, signatures are not logged. Note: Model signatures are MLflow model attributes and are only collected if log_models is also True.

  • log_models – If True, trained models are logged as MLflow model artifacts. If False, trained models are not logged. Input examples and model signatures, which are attributes of MLflow models, are also omitted when log_models is False.

  • log_datasets – If True, train and validation dataset information is logged to MLflow Tracking if applicable. If False, dataset information is not logged.

  • disable – If True, disables the scikit-learn autologging integration. If False, enables the scikit-learn autologging integration.

  • exclusive – If True, autologged content is not logged to user-created fluent runs. If False, autologged content is logged to the active fluent run, which may be user-created.

  • disable_for_unsupported_versions – If True, disable autologging for versions of scikit-learn that have not been tested against this version of the MLflow client or are incompatible.

  • silent – If True, suppress all event logs and warnings from MLflow during scikit-learn autologging. If False, show all events and warnings during scikit-learn autologging.

  • max_tuning_runs – The maximum number of child MLflow runs created for hyperparameter search estimators. To create child runs for the best k results from the search, set max_tuning_runs to k. The default value is to track the best 5 search parameter sets. If max_tuning_runs=None, then a child run is created for each search parameter set. Note: The best k results is based on ordering in rank_test_score. In the case of multi-metric evaluation with a custom scorer, the first scorer’s rank_test_score_<scorer_name> will be used to select the best k results. To change metric used for selecting best k results, change ordering of dict passed as scoring parameter for estimator.

  • log_post_training_metrics – If True, post training metrics are logged. Defaults to True. See the post training metrics section for more details.

  • serialization_format – The format in which to serialize the model. This should be one of the following: mlflow.sklearn.SERIALIZATION_FORMAT_PICKLE or mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE.

  • registered_model_name – If given, each time a model is trained, it is registered as a new model version of the registered model with this name. The registered model is created if it does not already exist.

  • pos_label – If given, used as the positive label to compute binary classification training metrics such as precision, recall, f1, etc. This parameter should only be set for binary classification model. If used for multi-label model, the training metrics calculation will fail and the training metrics won’t be logged. If used for regression model, the parameter will be ignored.

  • extra_tags – A dictionary of extra tags to set on each managed run created by autologging.

mlflow.sklearn.get_default_conda_env(include_cloudpickle=False)[source]
Returns

The default Conda environment for MLflow Models produced by calls to save_model() and log_model().

mlflow.sklearn.get_default_pip_requirements(include_cloudpickle=False)[source]
Returns

A list of default pip requirements for MLflow Models produced by this flavor. Calls to save_model() and log_model() produce a pip environment that, at minimum, contains these requirements.

mlflow.sklearn.load_model(model_uri, dst_path=None)[source]

Load a scikit-learn model from a local file or a run.

Parameters

  • model_uri –

    The location, in URI format, of the MLflow model, for example:

    • /Users/me/path/to/local/model

    • relative/path/to/local/model

    • s3://my_bucket/path/to/model

    • runs:/<mlflow_run_id>/run-relative/path/to/model

    • models:/<model_name>/<model_version>

    • models:/<model_name>/<stage>

    For more information about supported URI schemes, see Referencing Artifacts.

  • dst_path – The local filesystem path to which to download the model artifact. This directory must already exist. If unspecified, a local output path will be created.

Returns

A scikit-learn model.

Example 
import mlflow.sklearn

sk_model = mlflow.sklearn.load_model("runs:/96771d893a5e46159d9f3b49bf9013e2/sk_models")

# use Pandas DataFrame to make predictions
pandas_df = ...
predictions = sk_model.predict(pandas_df)
mlflow.sklearn.log_model(sk_model, artifact_path, conda_env=None, code_paths=None, serialization_format='cloudpickle', registered_model_name=None, signature: mlflow.models.signature.ModelSignature = None, input_example: Union[pandas.core.frame.DataFrame, numpy.ndarray, dict, list, csr_matrix, csc_matrix, str, bytes, tuple] = None, await_registration_for=300, pip_requirements=None, extra_pip_requirements=None, pyfunc_predict_fn='predict', metadata=None)[source]

Log a scikit-learn model as an MLflow artifact for the current run. Produces an MLflow Model containing the following flavors:

  • mlflow.sklearn

  • mlflow.pyfunc. NOTE: This flavor is only included for scikit-learn models that define predict(), since predict() is required for pyfunc model inference.

Parameters

  • sk_model – scikit-learn model to be saved.

  • artifact_path – Run-relative artifact path.

  • conda_env –

    Either a dictionary representation of a Conda environment or the path to a conda environment yaml file. If provided, this describes the environment this model should be run in. At a minimum, it should specify the dependencies contained in get_default_conda_env(). If None, a conda environment with pip requirements inferred by mlflow.models.infer_pip_requirements() is added to the model. If the requirement inference fails, it falls back to using get_default_pip_requirements(). pip requirements from conda_env are written to a pip requirements.txt file and the full conda environment is written to conda.yaml. The following is an example dictionary representation of a conda environment:

    {
    "name": "mlflow-env",
    "channels": ["conda-forge"],
    "dependencies": [
        "python=3.8.15",
        {
            "pip": [
                "scikit-learn==x.y.z"
            ],
        },
    ],
}

  • code_paths –

    A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files are prepended to the system path when the model is loaded. Files declared as dependencies for a given model should have relative imports declared from a common root path if multiple files are defined with import dependencies between them to avoid import errors when loading the model.

    For a detailed explanation of code_paths functionality, recommended usage patterns and limitations, see the code_paths usage guide.

  • serialization_format – The format in which to serialize the model. This should be one of the formats listed in mlflow.sklearn.SUPPORTED_SERIALIZATION_FORMATS. The Cloudpickle format, mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE, provides better cross-system compatibility by identifying and packaging code dependencies with the serialized model.

  • registered_model_name – If given, create a model version under registered_model_name, also creating a registered model if one with the given name does not exist.

  • signature –

    an instance of the ModelSignature class that describes the model’s inputs and outputs. If not specified but an input_example is supplied, a signature will be automatically inferred based on the supplied input example and model. To disable automatic signature inference when providing an input example, set signature to False. To manually infer a model signature, call infer_signature() on datasets with valid model inputs, such as a training dataset with the target column omitted, and valid model outputs, like model predictions made on the training dataset, for example:

    from mlflow.models import infer_signature

train = df.drop_column("target_label")
predictions = ...  # compute model predictions
signature = infer_signature(train, predictions)

  • input_example – one or several instances of valid model input. The input example is used as a hint of what data to feed the model. It will be converted to a Pandas DataFrame and then serialized to json using the Pandas split-oriented format, or a numpy array where the example will be serialized to json by converting it to a list. Bytes are base64-encoded. When the signature parameter is None, the input example is used to infer a model signature.

  • await_registration_for – Number of seconds to wait for the model version to finish being created and is in READY status. By default, the function waits for five minutes. Specify 0 or None to skip waiting.

  • pip_requirements – Either an iterable of pip requirement strings (e.g. ["scikit-learn", "-r requirements.txt", "-c constraints.txt"]) or the string path to a pip requirements file on the local filesystem (e.g. "requirements.txt"). If provided, this describes the environment this model should be run in. If None, a default list of requirements is inferred by mlflow.models.infer_pip_requirements() from the current software environment. If the requirement inference fails, it falls back to using get_default_pip_requirements(). Both requirements and constraints are automatically parsed and written to requirements.txt and constraints.txt files, respectively, and stored as part of the model. Requirements are also written to the pip section of the model’s conda environment (conda.yaml) file.

  • extra_pip_requirements –

    Either an iterable of pip requirement strings (e.g. ["pandas", "-r requirements.txt", "-c constraints.txt"]) or the string path to a pip requirements file on the local filesystem (e.g. "requirements.txt"). If provided, this describes additional pip requirements that are appended to a default set of pip requirements generated automatically based on the user’s current software environment. Both requirements and constraints are automatically parsed and written to requirements.txt and constraints.txt files, respectively, and stored as part of the model. Requirements are also written to the pip section of the model’s conda environment (conda.yaml) file.

    Warning

    The following arguments can’t be specified at the same time:

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    This example demonstrates how to specify pip requirements using pip_requirements and extra_pip_requirements.

  • pyfunc_predict_fn – The name of the prediction function to use for inference with the pyfunc representation of the resulting MLflow Model. Current supported functions are: "predict", "predict_proba", "predict_log_proba", "predict_joint_log_proba", and "score".

  • metadata – Custom metadata dictionary passed to the model and stored in the MLmodel file.

Returns

A ModelInfo instance that contains the metadata of the logged model.

Example 
import mlflow
import mlflow.sklearn
from mlflow.models import infer_signature
from sklearn.datasets import load_iris
from sklearn import tree

with mlflow.start_run():
    # load dataset and train model
    iris = load_iris()
    sk_model = tree.DecisionTreeClassifier()
    sk_model = sk_model.fit(iris.data, iris.target)

    # log model params
    mlflow.log_param("criterion", sk_model.criterion)
    mlflow.log_param("splitter", sk_model.splitter)
    signature = infer_signature(iris.data, sk_model.predict(iris.data))

    # log model
    mlflow.sklearn.log_model(sk_model, "sk_models", signature=signature)
mlflow.sklearn.save_model(sk_model, path, conda_env=None, code_paths=None, mlflow_model=None, serialization_format='cloudpickle', signature: mlflow.models.signature.ModelSignature = None, input_example: Union[pandas.core.frame.DataFrame, numpy.ndarray, dict, list, csr_matrix, csc_matrix, str, bytes, tuple] = None, pip_requirements=None, extra_pip_requirements=None, pyfunc_predict_fn='predict', metadata=None)[source]

Save a scikit-learn model to a path on the local file system. Produces a MLflow Model containing the following flavors:

  • mlflow.sklearn

  • mlflow.pyfunc. NOTE: This flavor is only included for scikit-learn models that define predict(), since predict() is required for pyfunc model inference.

Parameters

  • sk_model – scikit-learn model to be saved.

  • path – Local path where the model is to be saved.

  • conda_env –

    Either a dictionary representation of a Conda environment or the path to a conda environment yaml file. If provided, this describes the environment this model should be run in. At a minimum, it should specify the dependencies contained in get_default_conda_env(). If None, a conda environment with pip requirements inferred by mlflow.models.infer_pip_requirements() is added to the model. If the requirement inference fails, it falls back to using get_default_pip_requirements(). pip requirements from conda_env are written to a pip requirements.txt file and the full conda environment is written to conda.yaml. The following is an example dictionary representation of a conda environment:

    {
    "name": "mlflow-env",
    "channels": ["conda-forge"],
    "dependencies": [
        "python=3.8.15",
        {
            "pip": [
                "scikit-learn==x.y.z"
            ],
        },
    ],
}

  • code_paths –

    A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files are prepended to the system path when the model is loaded. Files declared as dependencies for a given model should have relative imports declared from a common root path if multiple files are defined with import dependencies between them to avoid import errors when loading the model.

    For a detailed explanation of code_paths functionality, recommended usage patterns and limitations, see the code_paths usage guide.

  • mlflow_model – mlflow.models.Model this flavor is being added to.

  • serialization_format – The format in which to serialize the model. This should be one of the formats listed in mlflow.sklearn.SUPPORTED_SERIALIZATION_FORMATS. The Cloudpickle format, mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE, provides better cross-system compatibility by identifying and packaging code dependencies with the serialized model.

  • signature –

    an instance of the ModelSignature class that describes the model’s inputs and outputs. If not specified but an input_example is supplied, a signature will be automatically inferred based on the supplied input example and model. To disable automatic signature inference when providing an input example, set signature to False. To manually infer a model signature, call infer_signature() on datasets with valid model inputs, such as a training dataset with the target column omitted, and valid model outputs, like model predictions made on the training dataset, for example:

    from mlflow.models import infer_signature

train = df.drop_column("target_label")
predictions = ...  # compute model predictions
signature = infer_signature(train, predictions)

  • input_example – one or several instances of valid model input. The input example is used as a hint of what data to feed the model. It will be converted to a Pandas DataFrame and then serialized to json using the Pandas split-oriented format, or a numpy array where the example will be serialized to json by converting it to a list. Bytes are base64-encoded. When the signature parameter is None, the input example is used to infer a model signature.

  • pip_requirements – Either an iterable of pip requirement strings (e.g. ["scikit-learn", "-r requirements.txt", "-c constraints.txt"]) or the string path to a pip requirements file on the local filesystem (e.g. "requirements.txt"). If provided, this describes the environment this model should be run in. If None, a default list of requirements is inferred by mlflow.models.infer_pip_requirements() from the current software environment. If the requirement inference fails, it falls back to using get_default_pip_requirements(). Both requirements and constraints are automatically parsed and written to requirements.txt and constraints.txt files, respectively, and stored as part of the model. Requirements are also written to the pip section of the model’s conda environment (conda.yaml) file.

  • extra_pip_requirements –

    Either an iterable of pip requirement strings (e.g. ["pandas", "-r requirements.txt", "-c constraints.txt"]) or the string path to a pip requirements file on the local filesystem (e.g. "requirements.txt"). If provided, this describes additional pip requirements that are appended to a default set of pip requirements generated automatically based on the user’s current software environment. Both requirements and constraints are automatically parsed and written to requirements.txt and constraints.txt files, respectively, and stored as part of the model. Requirements are also written to the pip section of the model’s conda environment (conda.yaml) file.

    Warning

    The following arguments can’t be specified at the same time:

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    This example demonstrates how to specify pip requirements using pip_requirements and extra_pip_requirements.

  • pyfunc_predict_fn – The name of the prediction function to use for inference with the pyfunc representation of the resulting MLflow Model. Current supported functions are: "predict", "predict_proba", "predict_log_proba", "predict_joint_log_proba", and "score".

  • metadata – Custom metadata dictionary passed to the model and stored in the MLmodel file.

Example 
import mlflow.sklearn
from sklearn.datasets import load_iris
from sklearn import tree

iris = load_iris()
sk_model = tree.DecisionTreeClassifier()
sk_model = sk_model.fit(iris.data, iris.target)

# Save the model in cloudpickle format
# set path to location for persistence
sk_path_dir_1 = ...
mlflow.sklearn.save_model(
    sk_model,
    sk_path_dir_1,
    serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE,
)

# save the model in pickle format
# set path to location for persistence
sk_path_dir_2 = ...
mlflow.sklearn.save_model(
    sk_model,
    sk_path_dir_2,
    serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_PICKLE,
)