Weights & Biases Reference Guide

Weights & Biases (W&B) is a platform for ML experiment tracking and MLOps. It enables automatic metric logging, real-time visualization, hyperparameter optimization, model registry, artifact versioning, and team collaboration.

When to Use W&B

Use W&B to:

• Track ML experiments with automatic metric logging
• Visualize training in real-time dashboards
• Compare runs across hyperparameters and configurations
• Optimize hyperparameters with automated sweeps
• Manage model registry with versioning and lineage
• Collaborate on ML projects with team workspaces
• Track artifacts (datasets, models, code) with lineage

Prerequisites

• Python environment with pip
• W&B account (free at wandb.ai)
• API key (obtained via wandb login or from W&B settings)
• For integrations: respective frameworks installed (PyTorch, TensorFlow, HuggingFace, etc.)

Installation and Login

Install the package and authenticate.

# Install W&B
pip install wandb

# Login (creates API key)
wandb login

# Or set API key programmatically
export WANDB_API_KEY=your_api_key_here

The API key must be kept secret; do not commit it to version control. The free tier includes unlimited public projects and 100GB storage; private projects require Teams or Enterprise. A team account on wandb.ai is required for team collaboration.

Basic Experiment Tracking

Start a run, log metrics, and finish.

import wandb

# Initialize a run
run = wandb.init(
    project="my-project",
    config={
        "learning_rate": 0.001,
        "epochs": 10,
        "batch_size": 32,
        "architecture": "ResNet50"
    }
)

# Training loop
for epoch in range(run.config.epochs):
    # Your training code
    train_loss = train_epoch()
    val_loss = validate()

    # Log metrics
    wandb.log({
        "epoch": epoch,
        "train/loss": train_loss,
        "val/loss": val_loss,
        "train/accuracy": train_acc,
        "val/accuracy": val_acc
    })

# Finish the run
wandb.finish()

Parameters for wandb.init

Parameter	Meaning	Required
project	Name of the project (collection of related experiments)	Yes
config	Dictionary of hyperparameters and configuration values	No (recommended)
name	Optional human-readable run name	No
tags	List of tags for organizing runs	No
notes	Free-text notes for the run	No
group	Group name to relate multiple runs	No
job_type	Type of job (e.g., 'train', 'eval')	No

Logging Metrics

Use wandb.log() to record metrics during training.

# Log scalars
wandb.log({"loss": 0.5, "accuracy": 0.92})

# Log multiple metrics
wandb.log({
    "train/loss": train_loss,
    "train/accuracy": train_acc,
    "val/loss": val_loss,
    "val/accuracy": val_acc,
    "learning_rate": current_lr,
    "epoch": epoch
})

# Log with custom x-axis
wandb.log({"loss": loss}, step=global_step)

# Log media (images, audio, video)
wandb.log({"examples": [wandb.Image(img) for img in images]})

# Log histograms
wandb.log({"gradients": wandb.Histogram(gradients)})

# Log tables
table = wandb.Table(columns=["id", "prediction", "ground_truth"])
wandb.log({"predictions": table})

The step parameter sets a custom x-axis value. Run names should be descriptive for easy identification.

# ✅ Good: Descriptive run names
wandb.init(
    project="nlp-classification",
    name="bert-base-lr0.001-bs32-epoch10"
)

# ❌ Bad: Generic names
wandb.init(project="nlp", name="run1")

Run Organization

wandb.init(
    project="my-project",
    tags=["baseline", "resnet50", "imagenet"],
    group="resnet-experiments",  # Group related runs
    job_type="train"             # Type of job
)

Logging System and Metadata

# Log system metrics
wandb.log({
    "gpu/util": gpu_utilization,
    "gpu/memory": gpu_memory_used,
    "cpu/util": cpu_utilization
})

# Log code version
wandb.log({"git_commit": git_commit_hash})

# Log data splits
wandb.log({
    "data/train_size": len(train_dataset),
    "data/val_size": len(val_dataset)
})

Saving Models and Checkpoints

import torch
import wandb

# Save model checkpoint
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}

torch.save(checkpoint, 'checkpoint.pth')

# Upload to W&B
wandb.save('checkpoint.pth')

# Or use Artifacts (recommended)
artifact = wandb.Artifact('model', type='model')
artifact.add_file('checkpoint.pth')
wandb.log_artifact(artifact)

Final Artifact and Predictions

# Save final model
artifact = wandb.Artifact('final-model', type='model')
artifact.add_file('model.pth')
wandb.log_artifact(artifact)

# Save predictions for analysis
predictions_table = wandb.Table(
    columns=["id", "input", "prediction", "ground_truth"],
    data=predictions_data
)
wandb.log({"predictions": predictions_table})

Run Sharing

# Runs are automatically shareable via URL
run = wandb.init(project="team-project")
print(f"Share this URL: {run.url}")

PyTorch Integration

import torch
import wandb

# Initialize
wandb.init(project="pytorch-demo", config={
    "lr": 0.001,
    "epochs": 10
})

# Access config
config = wandb.config

# Training loop
for epoch in range(config.epochs):
    for batch_idx, (data, target) in enumerate(train_loader):
        # Forward pass
        output = model(data)
        loss = criterion(output, target)

        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Log every 100 batches
        if batch_idx % 100 == 0:
            wandb.log({
                "loss": loss.item(),
                "epoch": epoch,
                "batch": batch_idx
            })

# Save model
torch.save(model.state_dict(), "model.pth")
wandb.save("model.pth")  # Upload to W&B

wandb.finish()

Hyperparameter Sweeps

Define a sweep configuration, create the sweep, and run an agent.

sweep_config = {
    'method': 'bayes',  # or 'grid', 'random'
    'metric': {
        'name': 'val/accuracy',
        'goal': 'maximize'
    },
    'parameters': {
        'learning_rate': {
            'distribution': 'log_uniform_values',
            'min': 1e-5,
            'max': 1e-1
        },
        'batch_size': {
            'values': [16, 32, 64, 128]
        },
        'optimizer': {
            'values': ['adam', 'sgd', 'rmsprop']
        },
        'dropout': {
            'distribution': 'uniform',
            'min': 0.1,
            'max': 0.5
        }
    }
}

# Initialize sweep
sweep_id = wandb.sweep(sweep_config, project="my-project")

Sweep Agent

def train():
    # Initialize run
    run = wandb.init()

    # Access sweep parameters
    lr = wandb.config.learning_rate
    batch_size = wandb.config.batch_size
    optimizer_name = wandb.config.optimizer

    # Build model with sweep config
    model = build_model(wandb.config)
    optimizer = get_optimizer(optimizer_name, lr)

    # Training loop
    for epoch in range(NUM_EPOCHS):
        train_loss = train_epoch(model, optimizer, batch_size)
        val_acc = validate(model)

        # Log metrics
        wandb.log({
            "train/loss": train_loss,
            "val/accuracy": val_acc
        })

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

The sweep agent runs indefinitely if count is not specified. Sweep configuration errors (e.g., missing metric for bayes) cause agent failure.

Sweep Methods

# Grid search - exhaustive
sweep_config = {
    'method': 'grid',
    'parameters': {
        'lr': {'values': [0.001, 0.01, 0.1]},
        'batch_size': {'values': [16, 32, 64]}
    }
}

# Random search
sweep_config = {
    'method': 'random',
    'parameters': {
        'lr': {'distribution': 'uniform', 'min': 0.0001, 'max': 0.1},
        'dropout': {'distribution': 'uniform', 'min': 0.1, 'max': 0.5}
    }
}

# Bayesian optimization (recommended)
sweep_config = {
    'method': 'bayes',
    'metric': {'name': 'val/loss', 'goal': 'minimize'},
    'parameters': {
        'lr': {'distribution': 'log_uniform_values', 'min': 1e-5, 'max': 1e-1}
    }
}

Sweep Parameters

Parameter	Meaning	Required
method	Sweep strategy: 'grid', 'random', or 'bayes'	Yes (in sweep_config)
metric	Dictionary with 'name' (metric to optimize) and 'goal' ('minimize' or 'maximize')	Yes (for bayes and random sweeps)
parameters	Dictionary of hyperparameter distributions and ranges	Yes (in sweep_config)
count	Number of trials to run in wandb.agent	No (default runs indefinitely)

Artifact Tracking

Artifacts track datasets, models, and code with lineage.

Logging Artifacts

# Create artifact
artifact = wandb.Artifact(
    name='training-dataset',
    type='dataset',
    description='ImageNet training split',
    metadata={'size': '1.2M images', 'split': 'train'}
)

# Add files
artifact.add_file('data/train.csv')
artifact.add_dir('data/images/')

# Log artifact
wandb.log_artifact(artifact)

Using Artifacts

# Download and use artifact
run = wandb.init(project="my-project")

# Download artifact
artifact = run.use_artifact('training-dataset:latest')
artifact_dir = artifact.download()

# Use the data
data = load_data(f"{artifact_dir}/train.csv")

Artifact names must be unique within a project; use aliases for versioning. Name conflicts overwrite previous versions unless aliases are used. Exceeding storage limits on the free tier blocks artifact logging.

Model Registry

Log model artifacts with aliases and link them to a registry.

# Log model as artifact
model_artifact = wandb.Artifact(
    name='resnet50-model',
    type='model',
    metadata={'architecture': 'ResNet50', 'accuracy': 0.95}
)

model_artifact.add_file('model.pth')
wandb.log_artifact(model_artifact, aliases=['best', 'production'])

# Link to model registry
run.link_artifact(model_artifact, 'model-registry/production-models')

The aliases parameter is a list of aliases for artifact versioning (e.g., 'best', 'production').

HuggingFace Transformers Integration

from transformers import Trainer, TrainingArguments
import wandb

# Initialize W&B
wandb.init(project="hf-transformers")

# Training arguments with W&B
training_args = TrainingArguments(
    output_dir="./results",
    report_to="wandb",  # Enable W&B logging
    run_name="bert-finetuning",
    logging_steps=100,
    save_steps=500
)

# Trainer automatically logs to W&B
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset
)

trainer.train()

Set report_to to 'wandb' in HuggingFace TrainingArguments to enable logging.

PyTorch Lightning Integration

from pytorch_lightning import Trainer
from pytorch_lightning.loggers import WandbLogger
import wandb

# Create W&B logger
wandb_logger = WandbLogger(
    project="lightning-demo",
    log_model=True  # Log model checkpoints
)

# Use with Trainer
trainer = Trainer(
    logger=wandb_logger,
    max_epochs=10
)

trainer.fit(model, datamodule=dm)

The log_model parameter in WandbLogger is a boolean to log model checkpoints.

Keras/TensorFlow Integration

import wandb
from wandb.integration.keras import WandbMetricsLogger, WandbModelCheckpoint

# Initialize
wandb.init(project="keras-demo")

# Add callbacks (the monolithic WandbCallback was removed;
# use the dedicated callbacks from wandb.integration.keras instead)
model.fit(
    x_train, y_train,
    validation_data=(x_val, y_val),
    epochs=10,
    callbacks=[
        WandbMetricsLogger(),                        # Auto-logs metrics
        WandbModelCheckpoint("models/model-{epoch}")  # Saves checkpoints
    ]
)

The monolithic WandbCallback for Keras was removed; use dedicated callbacks from wandb.integration.keras instead. Missing or incorrect callbacks cause errors.

Custom Visualizations

# Log custom visualizations
import matplotlib.pyplot as plt

fig, ax = plt.subplots()
ax.plot(x, y)
wandb.log({"custom_plot": wandb.Image(fig)})

# Log confusion matrix
wandb.log({"conf_mat": wandb.plot.confusion_matrix(
    probs=None,
    y_true=ground_truth,
    preds=predictions,
    class_names=class_names
)})

Offline Mode

Use offline mode when network connectivity is unstable.

import os

# Enable offline mode
os.environ["WANDB_MODE"] = "offline"

wandb.init(project="my-project")
# ... your code ...

# Sync later
# wandb sync <run_directory>

Set the environment variable WANDB_MODE to 'offline' to disable network sync. After runs complete, sync manually with wandb sync. Setting WANDB_MODE incorrectly may disable logging unexpectedly.

Failure Modes

• Network issues cause logging failures; use offline mode as workaround.
• Incorrect API key or missing login prevents logging.
• Sweep configuration errors (e.g., missing metric for bayes) cause agent failure.
• Artifact name conflicts overwrite previous versions unless aliases are used.
• Exceeding storage limits on free tier blocks artifact logging.
• Missing or incorrect callbacks in Keras integration (using old WandbCallback) causes errors.
• Run not finished with wandb.finish() may cause incomplete data.
• Environment variable WANDB_MODE set incorrectly may disable logging unexpectedly.

Additional Resources

For detailed documentation on specific topics, refer to references/sweeps.md, references/artifacts.md, and references/integrations.md.

Configuration Example

# Create/use project
run = wandb.init(
    project="image-classification",
    name="resnet50-experiment-1",  # Optional run name
    tags=["baseline", "resnet"],    # Organize with tags
    notes="First baseline run"      # Add notes
)

# Each run has unique ID
print(f"Run ID: {run.id}")
print(f"Run URL: {run.url}")

Config Dictionary

config = {
    # Model architecture
    "model": "ResNet50",
    "pretrained": True,

    # Training params
    "learning_rate": 0.001,
    "batch_size": 32,
    "epochs": 50,
    "optimizer": "Adam",

    # Data params
    "dataset": "ImageNet",
    "augmentation": "standard"
}

wandb.init(project="my-project", config=config)

# Access config during training
lr = wandb.config.learning_rate
batch_size = wandb.config.batch_size

Capabilities Summary

• Automatic metric logging and real-time dashboards
• Run comparison across hyperparameters and configurations
• Hyperparameter sweeps (grid, random, Bayesian)
• Model registry with versioning and lineage
• Artifact tracking for datasets, models, and code
• Team workspaces and collaboration
• Integration with 100+ tools (PyTorch, HuggingFace, Lightning, Keras, etc.)
• Custom visualizations (images, histograms, tables, confusion matrices)
• Shareable reports with markdown and embedded visualizations
• Offline mode for unstable connections
• Model checkpointing and artifact logging

The platform costs $50 per month for the Pro tier, and supports up to 200,000 runs per project on the Enterprise plan.