Chapter 2: Weights and Biases

Introduction to W&B

Weights & Biases (W&B) is a powerful MLOps platform for experiment tracking, dataset versioning, and model management with excellent visualization capabilities.

Setup and Configuration

# Install
pip install wandb

# Login
wandb login

# Or set API key
export WANDB_API_KEY="your-api-key"

Basic Experiment Tracking

import wandb
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, roc_auc_score

# Initialize a run
wandb.init(
    project="customer-churn",
    name="rf-baseline",
    config={
        "n_estimators": 100,
        "max_depth": 10,
        "learning_rate": 0.1,
        "dataset": "v2.0"
    }
)

# Access config
config = wandb.config

# Train model
model = RandomForestClassifier(
    n_estimators=config.n_estimators,
    max_depth=config.max_depth
)
model.fit(X_train, y_train)

# Log metrics
accuracy = accuracy_score(y_test, model.predict(X_test))
auc = roc_auc_score(y_test, model.predict_proba(X_test)[:, 1])

wandb.log({
    "accuracy": accuracy,
    "auc_roc": auc,
    "epoch": 1
})

# Finish run
wandb.finish()

Hyperparameter Sweeps

# Define sweep configuration
sweep_config = {
    "method": "bayes",  # grid, random, bayes
    "metric": {"name": "auc_roc", "goal": "maximize"},
    "parameters": {
        "n_estimators": {"values": [100, 200, 300]},
        "max_depth": {"min": 5, "max": 20},
        "min_samples_split": {"values": [2, 5, 10]},
        "learning_rate": {"distribution": "log_uniform_values", "min": 0.001, "max": 0.1}
    }
}

# Create sweep
sweep_id = wandb.sweep(sweep_config, project="customer-churn")

# Training function
def train():
    wandb.init()
    config = wandb.config
    
    model = GradientBoostingClassifier(
        n_estimators=config.n_estimators,
        max_depth=config.max_depth,
        learning_rate=config.learning_rate
    )
    model.fit(X_train, y_train)
    
    auc = roc_auc_score(y_test, model.predict_proba(X_test)[:, 1])
    wandb.log({"auc_roc": auc})

# Run sweep
wandb.agent(sweep_id, train, count=50)

W&B Artifacts

# Log dataset as artifact
artifact = wandb.Artifact("training-data", type="dataset")
artifact.add_file("data/train.csv")
artifact.add_dir("data/features/")
wandb.log_artifact(artifact)

# Log model as artifact
model_artifact = wandb.Artifact("churn-model", type="model")
model_artifact.add_file("models/model.pkl")
wandb.log_artifact(model_artifact)

# Use artifact in another run
wandb.init(project="customer-churn")
artifact = wandb.use_artifact("training-data:latest")
artifact_dir = artifact.download()

W&B Tables for Data Visualization

import wandb
import pandas as pd

wandb.init(project="customer-churn")

# Create a W&B Table
table = wandb.Table(dataframe=df.head(100))
wandb.log({"data_sample": table})

# Log predictions table
predictions_df = pd.DataFrame({
    "customer_id": test_ids,
    "actual": y_test,
    "predicted": predictions,
    "probability": probabilities
})
wandb.log({"predictions": wandb.Table(dataframe=predictions_df)})

Key Takeaways

Summary

1. W&B excels at visualization with interactive dashboards
2. Sweeps automate hyperparameter optimization with Bayesian methods
3. Artifacts track datasets and models with lineage
4. Tables enable rich data exploration in the UI
5. Reports facilitate team collaboration

---

Next: Chapter 3: Comparing Experiments