跳到主要内容

Scikit-learn与MLflow

在本综合指南中，我们将引导您了解如何将scikit-learn与MLflow结合使用，以进行实验跟踪、模型管理和生产部署。我们将涵盖自动日志记录和手动日志记录方法，从基本用法到高级生产模式。

使用自动日志记录快速入门​

最快的入门方法是使用MLflow的scikit-learn自动日志记录。只需一行代码，您就可以自动跟踪scikit-learn实验中的参数、指标和模型。这种方法无需更改现有训练代码，并捕获了可复现ML工作流所需的一切。

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split

# Enable autologging for scikit-learn
mlflow.sklearn.autolog()

# Load sample data
wine = load_wine()
X_train, X_test, y_train, y_test = train_test_split(
    wine.data, wine.target, test_size=0.2, random_state=42
)

# Train your model - MLflow automatically logs everything
with mlflow.start_run():
    model = RandomForestClassifier(n_estimators=100, max_depth=3, random_state=42)
    model.fit(X_train, y_train)

    # Evaluation metrics are automatically captured
    train_score = model.score(X_train, y_train)
    test_score = model.score(X_test, y_test)

    print(f"Training accuracy: {train_score:.3f}")
    print(f"Test accuracy: {test_score:.3f}")

这个简单的示例会自动记录所有模型参数、训练指标、经过适当序列化的训练模型以及用于部署的模型签名——所有这些都无需任何额外代码。

理解自动日志记录行为​

MLflow的scikit-learn自动日志记录会自动捕获有关训练过程的全面信息。以下是每次训练模型时会跟踪的具体内容：

类别捕获的信息
参数来自estimator.get_params(deep=True)的所有参数
指标训练分数、分类/回归指标
标签估计器类名和完全限定类名
工件序列化模型、模型签名、指标信息

自动日志记录系统旨在全面且非侵入性。它捕获了可复现性所需的一切，而无需更改现有的scikit-learn代码。

日志记录方法​

为了完全控制日志记录的内容，您可以手动配置scikit-learn代码。当您需要自定义指标、特定工件日志记录或希望以特定方式组织实验时，这种方法是理想的选择。

import mlflow
import mlflow.sklearn
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from mlflow.models import infer_signature

# Generate sample data
X, y = make_classification(n_samples=1000, n_features=20, n_classes=2, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Manual logging approach
with mlflow.start_run():
    # Define hyperparameters
    params = {"C": 1.0, "max_iter": 1000, "solver": "lbfgs", "random_state": 42}

    # Log parameters
    mlflow.log_params(params)

    # Train model
    model = LogisticRegression(**params)
    model.fit(X_train, y_train)

    # Make predictions
    y_pred = model.predict(X_test)

    # Calculate and log metrics
    metrics = {
        "accuracy": accuracy_score(y_test, y_pred),
        "precision": precision_score(y_test, y_pred, average="weighted"),
        "recall": recall_score(y_test, y_pred, average="weighted"),
        "f1_score": f1_score(y_test, y_pred, average="weighted"),
    }
    mlflow.log_metrics(metrics)

    # Infer model signature
    signature = infer_signature(X_train, model.predict(X_train))

    # Log the model
    mlflow.sklearn.log_model(
        sk_model=model,
        name="model",
        signature=signature,
        input_example=X_train[:5],  # Sample input for documentation
    )

超参数调优​

MLflow为scikit-learn的超参数优化工具提供了出色的支持，自动为参数搜索实验创建有组织的子运行。这使得跟踪和比较不同的参数组合变得容易。

import mlflow
from sklearn.model_selection import GridSearchCV
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

# Enable autologging with hyperparameter tuning support
mlflow.sklearn.autolog(max_tuning_runs=10)  # Track top 10 parameter combinations

# Load data
digits = load_digits()
X_train, X_test, y_train, y_test = train_test_split(
    digits.data, digits.target, test_size=0.2, random_state=42
)

# Define parameter grid
param_grid = {
    "n_estimators": [50, 100, 200],
    "max_depth": [5, 10, 15, None],
    "min_samples_split": [2, 5, 10],
    "min_samples_leaf": [1, 2, 4],
}

with mlflow.start_run(run_name="Random Forest Hyperparameter Tuning"):
    # Create and fit GridSearchCV
    rf = RandomForestClassifier(random_state=42)
    grid_search = GridSearchCV(
        rf, param_grid, cv=5, scoring="accuracy", n_jobs=-1, verbose=1
    )

    grid_search.fit(X_train, y_train)

    # Best model evaluation
    best_score = grid_search.score(X_test, y_test)
    print(f"Best parameters: {grid_search.best_params_}")
    print(f"Best cross-validation score: {grid_search.best_score_:.3f}")
    print(f"Test score: {best_score:.3f}")

MLflow会自动创建一个包含整体搜索结果的父运行，并为每个参数组合创建子运行，从而轻松分析哪些参数效果最好。

使用MLflow进行模型评估​

MLflow提供了一个全面的评估API，可以自动为scikit-learn模型生成指标、可视化和诊断工具。

import mlflow
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
from mlflow.models import infer_signature

# Load data and train model
wine = load_wine()
X_train, X_test, y_train, y_test = train_test_split(
    wine.data, wine.target, test_size=0.2, random_state=42
)

model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# Create evaluation dataset
eval_data = X_test.copy()
eval_data = pd.DataFrame(eval_data, columns=wine.feature_names)
eval_data["label"] = y_test

with mlflow.start_run():
    # Log model with signature
    signature = infer_signature(X_test, model.predict(X_test))
    mlflow.sklearn.log_model(model, name="model", signature=signature)
    model_uri = mlflow.get_artifact_uri("model")

    # Comprehensive evaluation with MLflow
    result = mlflow.evaluate(
        model_uri,
        eval_data,
        targets="label",
        model_type="classifier",  # or "regressor" for regression
        evaluators=["default"],
    )

    # Access automatic metrics
    print(f"Accuracy: {result.metrics['accuracy_score']:.3f}")
    print(f"F1 Score: {result.metrics['f1_score']:.3f}")
    print(f"ROC AUC: {result.metrics['roc_auc']:.3f}")

    # Access generated artifacts
    print("Generated artifacts:")
    for artifact_name, path in result.artifacts.items():
        print(f"  {artifact_name}: {path}")

自动生成包括

性能指标，如分类的准确率、精确率、召回率、F1分数、ROC-AUC。可视化，包括混淆矩阵、ROC曲线、精确率-召回率曲线。特征重要性，具有SHAP值和特征贡献分析。模型工件，所有图表和诊断信息都保存到MLflow。

模型比较和选择​

使用MLflow评估系统地比较多个scikit-learn模型。

from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC

# Define models to compare
sklearn_models = {
    "random_forest": RandomForestClassifier(n_estimators=100, random_state=42),
    "gradient_boosting": GradientBoostingClassifier(n_estimators=100, random_state=42),
    "logistic_regression": LogisticRegression(random_state=42, max_iter=1000),
    "svm": SVC(probability=True, random_state=42),
}

# Evaluate each model systematically
comparison_results = {}

for model_name, model in sklearn_models.items():
    with mlflow.start_run(run_name=f"eval_{model_name}"):
        # Train model
        model.fit(X_train, y_train)

        # Log model
        signature = infer_signature(X_train, model.predict(X_train))
        mlflow.sklearn.log_model(model, name="model", signature=signature)
        model_uri = mlflow.get_artifact_uri("model")

        # Comprehensive evaluation with MLflow
        result = mlflow.evaluate(
            model_uri,
            eval_data,
            targets="label",
            model_type="classifier",
            evaluators=["default"],
        )

        comparison_results[model_name] = result.metrics

        # Log key metrics for comparison
        mlflow.log_metrics(
            {
                "accuracy": result.metrics["accuracy_score"],
                "f1": result.metrics["f1_score"],
                "roc_auc": result.metrics["roc_auc"],
                "precision": result.metrics["precision_score"],
                "recall": result.metrics["recall_score"],
            }
        )

# Create comparison summary
import pandas as pd

comparison_df = pd.DataFrame(comparison_results).T
print("Model Comparison Summary:")
print(comparison_df[["accuracy_score", "f1_score", "roc_auc"]].round(3))

# Identify best model
best_model = comparison_df["f1_score"].idxmax()
print(f"\nBest model by F1 score: {best_model}")

模型验证和质量门​

使用MLflow的验证API确保scikit-learn模型质量。

from mlflow.models import MetricThreshold

# First, evaluate your scikit-learn model
result = mlflow.evaluate(model_uri, eval_data, targets="label", model_type="classifier")

# Define quality thresholds for classification models
quality_thresholds = {
    "accuracy_score": MetricThreshold(threshold=0.85, greater_is_better=True),
    "f1_score": MetricThreshold(threshold=0.80, greater_is_better=True),
    "roc_auc": MetricThreshold(threshold=0.75, greater_is_better=True),
}

# Validate model meets quality standards
try:
    mlflow.validate_evaluation_results(
        candidate_result=result,
        validation_thresholds=quality_thresholds,
    )
    print("✅ Scikit-learn model meets all quality thresholds")
except mlflow.exceptions.ModelValidationFailedException as e:
    print(f"❌ Model failed validation: {e}")

# Compare against baseline model (e.g., previous model version)
baseline_result = mlflow.evaluate(
    baseline_model_uri, eval_data, targets="label", model_type="classifier"
)

# Validate improvement over baseline
improvement_thresholds = {
    "f1_score": MetricThreshold(
        threshold=0.02, greater_is_better=True  # Must be 2% better
    ),
}

try:
    mlflow.validate_evaluation_results(
        candidate_result=result,
        baseline_result=baseline_result,
        validation_thresholds=improvement_thresholds,
    )
    print("✅ New model improves over baseline")
except mlflow.exceptions.ModelValidationFailedException as e:
    print(f"❌ Model doesn't improve sufficiently: {e}")

模型管理​

MLflow支持多种scikit-learn模型的序列化格式，每种格式都针对不同的部署场景进行了优化。了解这些选项有助于您为生产需求选择正确的方法。

import mlflow.sklearn
from sklearn.ensemble import RandomForestRegressor

model = RandomForestRegressor(n_estimators=100)
model.fit(X_train, y_train)

# Cloudpickle format (default) - better cross-system compatibility
mlflow.sklearn.log_model(
    sk_model=model,
    name="cloudpickle_model",
    serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE,
)

# Pickle format - faster but less portable
mlflow.sklearn.log_model(
    sk_model=model,
    name="pickle_model",
    serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_PICKLE,
)

Cloudpickle是默认格式，因为它通过识别并打包代码依赖项与序列化模型一起提供更好的跨系统兼容性。Pickle速度更快，但在不同环境中的可移植性较差。

生产部署​

模型注册表提供集中式模型管理，具有版本控制和基于别名的部署。这对于从开发到生产部署管理模型至关重要。

# Register model to MLflow Model Registry
import mlflow
from mlflow import MlflowClient

client = MlflowClient()

# Log and register model in one step
with mlflow.start_run():
    mlflow.sklearn.log_model(
        sk_model=model,
        name="model",
        registered_model_name="CustomerChurnModel",
        signature=signature,
    )

# Or register an existing model
run_id = "your_run_id"
model_uri = f"runs:/{run_id}/model"

# Register the model
registered_model = mlflow.register_model(model_uri=model_uri, name="CustomerChurnModel")

# Use aliases instead of deprecated stages for deployment management
# Set aliases for different deployment environments
client.set_registered_model_alias(
    name="CustomerChurnModel",
    alias="champion",  # Production model
    version=registered_model.version,
)

client.set_registered_model_alias(
    name="CustomerChurnModel",
    alias="challenger",  # A/B testing model
    version=registered_model.version,
)

# Use tags to track model status and metadata
client.set_model_version_tag(
    name="CustomerChurnModel",
    version=registered_model.version,
    key="validation_status",
    value="approved",
)

client.set_model_version_tag(
    name="CustomerChurnModel",
    version=registered_model.version,
    key="deployment_date",
    value="2025-05-29",
)

现代模型注册表功能

模型别名用灵活的命名引用替换了已弃用的阶段。您可以为任何模型版本分配多个别名（例如，champion、challenger、shadow），独立于模型训练更新别名以实现无缝部署，并将其用于A/B测试和渐进式发布。

模型标签提供丰富的元数据和状态跟踪。使用validation_status: approved跟踪验证状态，使用ready_for_prod: true标记部署准备情况，并使用team: data-science添加团队所有权。

基于环境的模型支持成熟的MLOps工作流。为每个环境创建单独的注册模型：dev.CustomerChurnModel、staging.CustomerChurnModel、prod.CustomerChurnModel，并使用copy_model_version()在环境之间提升模型。

# Promote model from staging to production environment
client.copy_model_version(
    src_model_uri="models:/staging.CustomerChurnModel@candidate",
    dst_name="prod.CustomerChurnModel",
)

高级功能​

Scikit-learn管道是MLflow中的一等公民，提供从数据预处理到模型训练的端到端工作流跟踪。这确保了整个ML工作流的可复现性。

import mlflow
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.feature_selection import SelectKBest, f_regression
from sklearn.linear_model import LinearRegression
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import OneHotEncoder

# Enable autologging for pipelines
mlflow.sklearn.autolog()

# Create a complex preprocessing and modeling pipeline
numeric_features = ["age", "income", "credit_score"]
categorical_features = ["occupation", "location"]

# Preprocessing pipeline
numeric_transformer = Pipeline(
    steps=[("scaler", StandardScaler()), ("selector", SelectKBest(f_regression, k=2))]
)

categorical_transformer = Pipeline(
    steps=[("encoder", OneHotEncoder(drop="first", sparse_output=False))]
)

# Combine preprocessing steps
preprocessor = ColumnTransformer(
    transformers=[
        ("num", numeric_transformer, numeric_features),
        ("cat", categorical_transformer, categorical_features),
    ]
)

# Complete pipeline with model
pipeline = Pipeline(
    steps=[("preprocessor", preprocessor), ("regressor", LinearRegression())]
)

# Train pipeline - all steps are automatically logged
with mlflow.start_run(run_name="Complete Pipeline Experiment"):
    pipeline.fit(X_train, y_train)

    # Pipeline scoring is automatically captured
    train_score = pipeline.score(X_train, y_train)
    test_score = pipeline.score(X_test, y_test)

    print(f"Pipeline R² score: {test_score:.3f}")

MLflow会自动记录每个管道阶段的参数，从而轻松了解数据是如何处理的以及使用了哪些模型参数。

结论​

MLflow的scikit-learn集成提供了用于传统机器学习工作流中实验跟踪、模型管理和部署的全面解决方案。无论您是使用简单的自动日志记录进行快速实验，还是实现复杂的生产管道，MLflow都能满足您的需求。

将MLflow与scikit-learn结合使用的主要优点

轻松的实验跟踪提供单行自动日志记录，捕获可复现ML所需的一切。超参数优化包括内置的网格搜索支持，具有有组织的子运行和轻松比较。全面评估通过mlflow.evaluate()提供自动指标生成、可视化和SHAP分析。生产就绪部署提供模型注册表集成，具有基于别名的部署和质量门。团队协作通过丰富的元数据和工件实现集中式实验管理。

本指南中的模式和示例为使用scikit-learn和MLflow构建可扩展、可复现的机器学习系统奠定了坚实的基础。从自动日志记录开始，立即获得好处，然后随着需求的增长逐步采用更高级的功能，如模型评估、注册表和自定义配置。