Introduction

Machine Learning development is iterative. A team may test different algorithms, feature sets, data versions, hyperparameters, random seeds, preprocessing strategies, decision thresholds, and runtime environments before selecting a model.

Without experiment tracking, results may be scattered across notebooks, terminal output, spreadsheets, screenshots, and filenames. It becomes difficult to determine which experiment produced a model, why one candidate was selected, or whether a historical result can be reproduced.

Experiment tracking is the systematic recording of the inputs, configuration, execution context, outputs, metrics, artifacts, and conclusions associated with each Machine Learning run.

MLflow and Weights & Biases, commonly abbreviated as W&B, provide experiment-tracking capabilities. Both can organize runs, log parameters and metrics, preserve artifacts, compare results, and improve collaboration, although their interfaces, deployment options, integrations, and broader platform capabilities differ.

Learning Objectives

Prerequisites

python -m pip install mlflow wandb pandas numpy scikit-learn matplotlib joblib

What Is a Machine Learning Experiment?

A Machine Learning experiment is a structured investigation that tests a defined hypothesis by changing one or more controlled inputs and measuring the resulting behavior.

What Is Experiment Tracking?

Experiment tracking creates a searchable history of model-development runs and connects every result to the inputs and artifacts that produced it.

Why Experiment Tracking Matters

Core Tracking Concepts

Parameters vs Metrics

Experiment Artifacts

Reproducibility Record

Create a Shared Training Function

The following training function is used in the MLflow and W&B examples so that both platforms track equivalent model behavior.

from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder
from sklearn.preprocessing import StandardScaler


def create_model_pipeline(
    numerical_features,
    categorical_features,
    regularization_strength,
    random_state
):
    numerical_pipeline = Pipeline(
        steps=[
            (
                "imputer",
                SimpleImputer(
                    strategy="median"
                )
            ),
            (
                "scaler",
                StandardScaler()
            )
        ]
    )

    categorical_pipeline = Pipeline(
        steps=[
            (
                "imputer",
                SimpleImputer(
                    strategy="most_frequent"
                )
            ),
            (
                "encoder",
                OneHotEncoder(
                    handle_unknown="ignore"
                )
            )
        ]
    )

    preprocessor = ColumnTransformer(
        transformers=[
            (
                "numerical",
                numerical_pipeline,
                numerical_features
            ),
            (
                "categorical",
                categorical_pipeline,
                categorical_features
            )
        ]
    )

    return Pipeline(
        steps=[
            (
                "preprocessor",
                preprocessor
            ),
            (
                "classifier",
                LogisticRegression(
                    C=regularization_strength,
                    class_weight="balanced",
                    max_iter=2000,
                    random_state=random_state
                )
            )
        ]
    )

Create an Evaluation Function

from sklearn.metrics import accuracy_score
from sklearn.metrics import average_precision_score
from sklearn.metrics import f1_score
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score
from sklearn.metrics import roc_auc_score


def calculate_metrics(
    true_labels,
    probabilities,
    decision_threshold
):
    predictions = (
        probabilities
        >=
        decision_threshold
    ).astype(
        int
    )

    return {
        "accuracy": float(
            accuracy_score(
                true_labels,
                predictions
            )
        ),
        "precision": float(
            precision_score(
                true_labels,
                predictions,
                zero_division=0
            )
        ),
        "recall": float(
            recall_score(
                true_labels,
                predictions,
                zero_division=0
            )
        ),
        "f1": float(
            f1_score(
                true_labels,
                predictions,
                zero_division=0
            )
        ),
        "roc_auc": float(
            roc_auc_score(
                true_labels,
                probabilities
            )
        ),
        "average_precision": float(
            average_precision_score(
                true_labels,
                probabilities
            )
        )
    }

MLflow Tracking Concepts

Configure MLflow Tracking

import mlflow


mlflow.set_tracking_uri(
    "http://127.0.0.1:5000"
)


mlflow.set_experiment(
    "customer-churn-experiments"
)

Start a Local MLflow Server

mlflow server \
  --host 127.0.0.1 \
  --port 5000

Track an Experiment with MLflow

import mlflow
import mlflow.sklearn


regularization_strength = 1.0

decision_threshold = 0.60

random_state = 42


with mlflow.start_run(
    run_name="logistic-regression-balanced"
) as run:

    mlflow.log_params(
        {
            "algorithm": (
                "logistic-regression"
            ),
            "regularization_strength": (
                regularization_strength
            ),
            "class_weight": "balanced",
            "decision_threshold": (
                decision_threshold
            ),
            "random_state": random_state
        }
    )

    mlflow.set_tags(
        {
            "project": "customer-churn",
            "task": "binary-classification",
            "code_version": (
                "replace-with-source-revision"
            ),
            "training_data_version": (
                "replace-with-data-version"
            ),
            "feature_definition_version": (
                "1.0.0"
            ),
            "run_purpose": (
                "evaluate-class-weighting"
            )
        }
    )

    model_pipeline = create_model_pipeline(
        numerical_features=(
            numerical_features
        ),
        categorical_features=(
            categorical_features
        ),
        regularization_strength=(
            regularization_strength
        ),
        random_state=random_state
    )

    model_pipeline.fit(
        training_features,
        training_labels
    )

    validation_probabilities = (
        model_pipeline.predict_proba(
            validation_features
        )[:, 1]
    )

    validation_metrics = (
        calculate_metrics(
            true_labels=validation_labels,
            probabilities=(
                validation_probabilities
            ),
            decision_threshold=(
                decision_threshold
            )
        )
    )

    mlflow.log_metrics(
        validation_metrics
    )

    mlflow.sklearn.log_model(
        sk_model=model_pipeline,
        name="model"
    )

    run_id = run.info.run_id


print(
    run_id
)

Log Evaluation Artifacts with MLflow

import json
from pathlib import Path


artifact_directory = Path(
    "experiment-artifacts"
)


artifact_directory.mkdir(
    parents=True,
    exist_ok=True
)


metrics_path = (
    artifact_directory
    / "validation_metrics.json"
)


with open(
    metrics_path,
    "w",
    encoding="utf-8"
) as metrics_file:

    json.dump(
        validation_metrics,
        metrics_file,
        indent=4
    )


mlflow.log_artifact(
    str(
        metrics_path
    ),
    artifact_path="evaluation"
)

Log an Input Example and Signature

from mlflow.models import infer_signature


input_example = (
    validation_features.head(
        3
    )
)


model_signature = infer_signature(
    input_example,
    model_pipeline.predict_proba(
        input_example
    )
)


mlflow.sklearn.log_model(
    sk_model=model_pipeline,
    name="model-with-signature",
    signature=model_signature,
    input_example=input_example
)

MLflow Autologging

MLflow autologging can automatically capture supported parameters, metrics, models, and related metadata from compatible libraries.

import mlflow
import mlflow.sklearn


mlflow.sklearn.autolog(
    log_input_examples=True,
    log_model_signatures=True,
    log_models=True
)


with mlflow.start_run(
    run_name="autologged-logistic-regression"
):

    model_pipeline = create_model_pipeline(
        numerical_features=(
            numerical_features
        ),
        categorical_features=(
            categorical_features
        ),
        regularization_strength=1.0,
        random_state=42
    )

    model_pipeline.fit(
        training_features,
        training_labels
    )

Log Metrics Over Time

with mlflow.start_run(
    run_name="iterative-training"
):

    for epoch in range(
        number_of_epochs
    ):
        training_loss = (
            calculate_training_loss(
                epoch
            )
        )

        validation_loss = (
            calculate_validation_loss(
                epoch
            )
        )

        mlflow.log_metric(
            "training_loss",
            float(
                training_loss
            ),
            step=epoch
        )

        mlflow.log_metric(
            "validation_loss",
            float(
                validation_loss
            ),
            step=epoch
        )

Search MLflow Runs

experiment = (
    mlflow.get_experiment_by_name(
        "customer-churn-experiments"
    )
)


runs = mlflow.search_runs(
    experiment_ids=[
        experiment.experiment_id
    ],
    filter_string=(
        "tags.task = "
        "'binary-classification'"
    ),
    order_by=[
        "metrics.f1 DESC"
    ],
    max_results=20
)


print(
    runs[
        [
            "run_id",
            "params.algorithm",
            "params.regularization_strength",
            "metrics.precision",
            "metrics.recall",
            "metrics.f1"
        ]
    ]
)

Load a Model from an MLflow Run

import mlflow.pyfunc


model_uri = (
    f"runs:/{run_id}/model"
)


tracked_model = mlflow.pyfunc.load_model(
    model_uri
)


predictions = tracked_model.predict(
    inference_frame
)

Weights & Biases Tracking Concepts

Authenticate Weights & Biases

wandb login

Track an Experiment with W&B

import wandb


configuration = {
    "algorithm": "logistic-regression",
    "regularization_strength": 1.0,
    "class_weight": "balanced",
    "decision_threshold": 0.60,
    "random_state": 42,
    "training_data_version": (
        "replace-with-data-version"
    ),
    "feature_definition_version": "1.0.0"
}


with wandb.init(
    entity="replace-with-approved-entity",
    project="customer-churn-experiments",
    name="logistic-regression-balanced",
    config=configuration,
    tags=[
        "binary-classification",
        "candidate",
        "scikit-learn"
    ],
    notes=(
        "Evaluate class weighting using "
        "the approved validation protocol."
    ),
    job_type="training"
) as run:

    model_pipeline = create_model_pipeline(
        numerical_features=(
            numerical_features
        ),
        categorical_features=(
            categorical_features
        ),
        regularization_strength=(
            run.config[
                "regularization_strength"
            ]
        ),
        random_state=(
            run.config[
                "random_state"
            ]
        )
    )

    model_pipeline.fit(
        training_features,
        training_labels
    )

    validation_probabilities = (
        model_pipeline.predict_proba(
            validation_features
        )[:, 1]
    )

    validation_metrics = (
        calculate_metrics(
            true_labels=validation_labels,
            probabilities=(
                validation_probabilities
            ),
            decision_threshold=(
                run.config[
                    "decision_threshold"
                ]
            )
        )
    )

    run.log(
        validation_metrics
    )

    run.summary[
        "code_version"
    ] = "replace-with-source-revision"

    run.summary[
        "validation_status"
    ] = "completed"

Log Training History with W&B

with wandb.init(
    entity="replace-with-approved-entity",
    project="neural-network-experiments",
    config={
        "epochs": number_of_epochs,
        "learning_rate": learning_rate
    }
) as run:

    for epoch in range(
        run.config[
            "epochs"
        ]
    ):
        training_loss = (
            calculate_training_loss(
                epoch
            )
        )

        validation_loss = (
            calculate_validation_loss(
                epoch
            )
        )

        run.log(
            {
                "epoch": epoch,
                "training_loss": float(
                    training_loss
                ),
                "validation_loss": float(
                    validation_loss
                )
            }
        )

Log a Model as a W&B Artifact

from pathlib import Path

import joblib
import wandb


artifact_directory = Path(
    "experiment-artifacts"
)


artifact_directory.mkdir(
    parents=True,
    exist_ok=True
)


model_path = (
    artifact_directory
    / "customer_churn_pipeline.joblib"
)


joblib.dump(
    model_pipeline,
    model_path
)


model_artifact = wandb.Artifact(
    name="customer-churn-classifier",
    type="model",
    description=(
        "Candidate customer-churn "
        "classification pipeline"
    ),
    metadata={
        "code_version": (
            "replace-with-source-revision"
        ),
        "training_data_version": (
            "replace-with-data-version"
        ),
        "decision_threshold": 0.60,
        "validation_metrics": (
            validation_metrics
        )
    }
)


model_artifact.add_file(
    str(
        model_path
    )
)


run.log_artifact(
    model_artifact
)

Track a Dataset as a W&B Artifact

dataset_artifact = wandb.Artifact(
    name="customer-churn-training-data",
    type="dataset",
    description=(
        "Approved training-data reference "
        "for the churn model"
    ),
    metadata={
        "schema_version": "1.0.0",
        "data_version": (
            "replace-with-data-version"
        ),
        "row_count": int(
            training_features.shape[0]
        ),
        "column_count": int(
            training_features.shape[1]
        )
    }
)


dataset_artifact.add_file(
    "data/training.parquet"
)


run.log_artifact(
    dataset_artifact
)

Log Prediction Samples with a W&B Table

prediction_table = wandb.Table(
    columns=[
        "record_id",
        "true_label",
        "probability",
        "predicted_label"
    ]
)


for (
    record_id,
    true_label,
    probability
) in zip(
    approved_record_identifiers,
    validation_labels,
    validation_probabilities
):
    predicted_label = int(
        probability
        >=
        decision_threshold
    )

    prediction_table.add_data(
        record_id,
        int(
            true_label
        ),
        float(
            probability
        ),
        predicted_label
    )


run.log(
    {
        "validation_predictions": (
            prediction_table
        )
    }
)

MLflow vs Weights & Biases

What Should You Compare Across Runs?

Compare Only Compatible Runs

Create a Run-Comparison Table

import pandas as pd


def create_run_comparison(
    run_records
):
    comparison = pd.DataFrame(
        run_records
    )

    required_columns = [
        "run_id",
        "algorithm",
        "data_version",
        "precision",
        "recall",
        "f1",
        "roc_auc",
        "decision_threshold"
    ]

    missing_columns = (
        set(required_columns)
        -
        set(comparison.columns)
    )

    if missing_columns:

        raise ValueError(
            "Missing comparison columns: "
            f"{sorted(missing_columns)}"
        )

    return comparison.sort_values(
        by=[
            "f1",
            "recall"
        ],
        ascending=[
            False,
            False
        ]
    )

Track a Confusion Matrix

import matplotlib.pyplot as plt
from sklearn.metrics import ConfusionMatrixDisplay


predicted_labels = (
    validation_probabilities
    >=
    decision_threshold
).astype(
    int
)


display = (
    ConfusionMatrixDisplay
    .from_predictions(
        validation_labels,
        predicted_labels
    )
)


figure_path = (
    artifact_directory
    / "confusion_matrix.png"
)


plt.savefig(
    figure_path,
    bbox_inches="tight"
)


plt.close()

mlflow.log_artifact(
    str(
        figure_path
    ),
    artifact_path="plots"
)

run.log(
    {
        "confusion_matrix": (
            wandb.Image(
                str(
                    figure_path
                )
            )
        )
    }
)

Track Threshold Selection

A classification threshold can materially change precision, recall, selection rate, and business impact. Track threshold evaluation separately from model training.

import pandas as pd


threshold_rows = []


for threshold in [
    0.30,
    0.40,
    0.50,
    0.60,
    0.70
]:
    metrics = calculate_metrics(
        true_labels=validation_labels,
        probabilities=(
            validation_probabilities
        ),
        decision_threshold=threshold
    )

    threshold_rows.append(
        {
            "threshold": threshold,
            **metrics
        }
    )


threshold_report = pd.DataFrame(
    threshold_rows
)


threshold_report.to_csv(
    artifact_directory
    / "threshold_analysis.csv",
    index=False
)

Track Cross-Validation Results

from sklearn.model_selection import cross_validate


cross_validation_results = cross_validate(
    estimator=model_pipeline,
    X=training_features,
    y=training_labels,
    cv=5,
    scoring=[
        "precision",
        "recall",
        "f1",
        "roc_auc"
    ],
    return_train_score=False
)


cross_validation_summary = {
    metric_name: float(
        values.mean()
    )
    for metric_name, values in (
        cross_validation_results.items()
    )
    if metric_name.startswith(
        "test_"
    )
}

mlflow.log_metrics(
    cross_validation_summary
)

run.log(
    cross_validation_summary
)

Hyperparameter Sweeps

A hyperparameter sweep coordinates multiple runs over a defined search space. Each trial should use the same approved evaluation protocol.

method: random

metric:
  name: validation_f1
  goal: maximize

parameters:
  regularization_strength:
    distribution: log_uniform_values
    min: 0.001
    max: 10.0

  class_weight:
    values:
      - balanced
      - null

  decision_threshold:
    values:
      - 0.4
      - 0.5
      - 0.6

W&B Sweep Training Function

import wandb


def train_sweep_candidate():

    with wandb.init() as run:

        model_pipeline = (
            create_model_pipeline(
                numerical_features=(
                    numerical_features
                ),
                categorical_features=(
                    categorical_features
                ),
                regularization_strength=(
                    run.config[
                        "regularization_strength"
                    ]
                ),
                random_state=42
            )
        )

        model_pipeline.set_params(
            classifier__class_weight=(
                run.config[
                    "class_weight"
                ]
            )
        )

        model_pipeline.fit(
            training_features,
            training_labels
        )

        probabilities = (
            model_pipeline.predict_proba(
                validation_features
            )[:, 1]
        )

        metrics = calculate_metrics(
            true_labels=validation_labels,
            probabilities=probabilities,
            decision_threshold=(
                run.config[
                    "decision_threshold"
                ]
            )
        )

        run.log(
            {
                f"validation_{key}": value
                for key, value in metrics.items()
            }
        )

Automatic Logging vs Explicit Logging

Naming and Tagging Strategy

Team Collaboration

Connect Runs to Source Control

git rev-parse HEAD

import subprocess


def get_git_commit():

    result = subprocess.run(
        [
            "git",
            "rev-parse",
            "HEAD"
        ],
        check=True,
        capture_output=True,
        text=True
    )

    return result.stdout.strip()


source_revision = get_git_commit()

mlflow.set_tag(
    "code_version",
    source_revision
)

run.summary[
    "code_version"
] = source_revision

Connect Runs to Data Versions

data_metadata = {
    "training_data_version": (
        "replace-with-training-version"
    ),
    "validation_data_version": (
        "replace-with-validation-version"
    ),
    "test_data_version": (
        "replace-with-test-version"
    ),
    "schema_version": "1.0.0",
    "feature_definition_version": "1.0.0"
}

mlflow.set_tags(
    data_metadata
)

for key, value in data_metadata.items():

    run.summary[
        key
    ] = value

Record Experiment Conclusions

{
    "hypothesis": "Class weighting may improve recall.",
    "controlled_change": "classifier class_weight",
    "evaluation_protocol": "approved validation split",
    "result_summary": "Populate from calculated experiment evidence.",
    "decision": "accept, reject, or investigate",
    "limitations": [
        "Populate observed limitations",
        "Do not infer from unmeasured behavior"
    ],
    "next_action": "Populate the approved next step"
}

Track Model-Selection Evidence

Integrate Tracking with an ML Pipeline

Integrate Tracking with CI/CD

ci_metadata = {
    "ci_build_id": (
        environment_build_id
    ),
    "source_revision": (
        source_revision
    ),
    "pipeline_version": "1.0.0",
    "release_status": "candidate"
}


mlflow.set_tags(
    ci_metadata
)

Security and Privacy

Tracking-Server Architecture

Retention and Run Hygiene

Responsible Experiment Tracking

Common Experiment-Tracking Mistakes

Recommended Experiment-Tracking Workflow

Hands-On Practice

Knowledge Check

Experiment-Tracking Readiness Checklist

Lesson Summary