使用 MLflow 跟踪实验：记录参数、指标与制品

Why Experiment Tracking Matters

Installing and Starting MLflow

# Install MLflow
# pip install mlflow scikit-learn

import mlflow
import mlflow.sklearn

# Start the tracking UI (run this in a terminal):
# mlflow ui --host 0.0.0.0 --port 5000
# Then open http://localhost:5000

# Check MLflow version
print('MLflow version:', mlflow.__version__)

# Default tracking URI stores to ./mlruns
print('Tracking URI:', mlflow.get_tracking_uri())

Creating Experiments

import mlflow

# Create or switch to an experiment
experiment_name = 'sentiment_classification'
mlflow.set_experiment(experiment_name)

# List all experiments
for exp in mlflow.search_experiments():
    print(f'ID: {exp.experiment_id}  Name: {exp.name}')

# Get current experiment info
experiment = mlflow.get_experiment_by_name(experiment_name)
print('Artifact location:', experiment.artifact_location)

Logging Parameters with mlflow.log_param

import mlflow
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

X, y = make_classification(n_samples=1000, n_features=20, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

with mlflow.start_run(run_name='rf_baseline'):
    n_estimators = 100
    max_depth = 5
    mlflow.log_param('n_estimators', n_estimators)
    mlflow.log_param('max_depth', max_depth)
    mlflow.log_param('random_state', 42)

    clf = RandomForestClassifier(n_estimators=n_estimators, max_depth=max_depth, random_state=42)
    clf.fit(X_train, y_train)
    print('Params logged.')

Logging Metrics with mlflow.log_metric

import mlflow
from sklearn.metrics import accuracy_score, f1_score

with mlflow.start_run(run_name='rf_run_1'):
    mlflow.log_param('n_estimators', 100)
    mlflow.log_param('max_depth', 5)

    # Train model (abbreviated)
    clf = RandomForestClassifier(n_estimators=100, max_depth=5, random_state=42)
    clf.fit(X_train, y_train)
    preds = clf.predict(X_test)

    acc = accuracy_score(y_test, preds)
    f1 = f1_score(y_test, preds)

    mlflow.log_metric('test_accuracy', acc)
    mlflow.log_metric('test_f1', f1)
    print(f'Logged accuracy={acc:.4f}, f1={f1:.4f}')

Logging Artifacts

import mlflow
import matplotlib.pyplot as plt
from sklearn.metrics import ConfusionMatrixDisplay
import numpy as np

with mlflow.start_run():
    clf.fit(X_train, y_train)
    preds = clf.predict(X_test)

    # Save confusion matrix as an artifact
    fig, ax = plt.subplots(figsize=(5, 4))
    ConfusionMatrixDisplay.from_predictions(y_test, preds, ax=ax)
    plt.tight_layout()
    plt.savefig('/tmp/confusion_matrix.png')
    plt.close()

    mlflow.log_artifact('/tmp/confusion_matrix.png')
    print('Artifact logged.')

Auto-logging with mlflow.sklearn.autolog

import mlflow
import mlflow.sklearn
from sklearn.ensemble import GradientBoostingClassifier

# Enable auto-logging -- no manual log calls needed
mlflow.sklearn.autolog()

with mlflow.start_run(run_name='gb_autolog'):
    clf = GradientBoostingClassifier(n_estimators=200, max_depth=3, learning_rate=0.05)
    clf.fit(X_train, y_train)
    # MLflow automatically logs:
    # - All constructor params
    # - Training accuracy
    # - Model artifact
    print('Auto-logged run complete.')

Comparing Runs in the MLflow UI

import mlflow

# Search runs programmatically (useful in CI/CD pipelines)
runs = mlflow.search_runs(
    experiment_names=['sentiment_classification'],
    order_by=['metrics.test_accuracy DESC']
)

if not runs.empty:
    best_run = runs.iloc[0]
    print('Best run ID:', best_run['run_id'])
    print('Best accuracy:', best_run['metrics.test_accuracy'])
    print('n_estimators:', best_run['params.n_estimators'])
    print('max_depth:', best_run['params.max_depth'])

Logging Models with mlflow.sklearn.log_model

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier

with mlflow.start_run(run_name='rf_logged_model'):
    clf = RandomForestClassifier(n_estimators=100, random_state=42)
    clf.fit(X_train, y_train)

    mlflow.log_metric('accuracy', accuracy_score(y_test, clf.predict(X_test)))

    # Log in sklearn flavour -- enables mlflow models serve
    mlflow.sklearn.log_model(
        sk_model=clf,
        artifact_path='random_forest',
        registered_model_name='SentimentRF'
    )
    print('Model registered in Model Registry.')

Tagging Runs for Organisation

import mlflow
import subprocess

def get_git_sha():
    try:
        return subprocess.check_output(
            ['git', 'rev-parse', 'HEAD'], text=True
        ).strip()
    except Exception:
        return 'unknown'

with mlflow.start_run():
    mlflow.set_tag('dataset_version', 'imdb_v2')
    mlflow.set_tag('git_sha', get_git_sha())
    mlflow.set_tag('author', 'mehmet.canker')
    mlflow.set_tag('notes', 'Testing higher max_depth after previous plateau')
    mlflow.log_param('max_depth', 8)
    print('Run tagged.')

Remote Tracking Server Setup

import mlflow
import os

# Point to a remote MLflow tracking server
os.environ['MLFLOW_TRACKING_URI'] = 'http://mlflow.yourcompany.com:5000'
os.environ['MLFLOW_TRACKING_USERNAME'] = 'data_team'
os.environ['MLFLOW_TRACKING_PASSWORD'] = 'secret'

mlflow.set_tracking_uri(os.environ['MLFLOW_TRACKING_URI'])

# All subsequent mlflow calls go to the remote server
mlflow.set_experiment('shared_team_experiment')
with mlflow.start_run():
    mlflow.log_param('model_type', 'xgboost')
    mlflow.log_metric('auc', 0.94)
    print('Logged to remote server.')

Quick Check

Lesson Recap