How to beginner · 3 min read

Bayesian optimization with wandb sweeps

Quick answer

Use wandb sweeps with the bayes search method to perform Bayesian optimization for hyperparameter tuning. Define a sweep configuration with parameter ranges and run your training script with wandb.agent to automatically explore the search space and optimize your model.

PREREQUISITES

Python 3.8+
wandb Python package (pip install wandb)
wandb account and API key
Basic knowledge of hyperparameter tuning

Setup

Install the wandb Python package and login to your Wandb account to enable sweep tracking.

Install wandb: pip install wandb
Login to wandb CLI: wandb login

bash

pip install wandb
wandb login

Step by step

Create a sweep configuration YAML or dictionary specifying method: bayes for Bayesian optimization, define the hyperparameters and their ranges, and the metric to optimize. Then initialize the sweep and run your training function with wandb.agent to perform the optimization.

python

import wandb
import random

def train():
    # Initialize a new wandb run
    with wandb.init() as run:
        # Access hyperparameters from sweep config
        lr = run.config.learning_rate
        batch_size = run.config.batch_size

        # Dummy training logic: simulate validation loss
        val_loss = (0.1 / lr) + (random.random() * 0.1) + (batch_size * 0.01)

        # Log the validation loss
        wandb.log({"val_loss": val_loss})

# Define sweep config for Bayesian optimization
sweep_config = {
    'method': 'bayes',  # Bayesian optimization
    'metric': {
        'name': 'val_loss',
        'goal': 'minimize'
    },
    'parameters': {
        'learning_rate': {
            'min': 0.0001,
            'max': 0.1
        },
        'batch_size': {
            'values': [16, 32, 64]
        }
    }
}

# Initialize sweep
sweep_id = wandb.sweep(sweep_config, project="bayes-opt-example")

# Run sweep agent to execute training with hyperparameter tuning
wandb.agent(sweep_id, function=train, count=10)

output

wandb: Creating sweep from: {'method': 'bayes', 'metric': {'name': 'val_loss', 'goal': 'minimize'}, 'parameters': {'learning_rate': {'min': 0.0001, 'max': 0.1}, 'batch_size': {'values': [16, 32, 64]}}}
wandb: Created sweep with ID: abc123
wandb: Starting sweep agent...
# Runs 10 training jobs with different hyperparameters, logging val_loss each run

Common variations

You can customize wandb sweeps by using different search methods like grid or random. For asynchronous or distributed tuning, run wandb.agent on multiple machines. You can also integrate with frameworks like PyTorch Lightning or TensorFlow by logging metrics inside their training loops.

python

import wandb

# Example sweep config with random search
random_sweep = {
    'method': 'random',
    'metric': {'name': 'val_loss', 'goal': 'minimize'},
    'parameters': {
        'learning_rate': {'min': 0.0001, 'max': 0.1},
        'batch_size': {'values': [16, 32, 64]}
    }
}

# Initialize and run sweep
sweep_id = wandb.sweep(random_sweep, project="random-search-example")
wandb.agent(sweep_id, function=train, count=5)

output

wandb: Created sweep with ID: def456
wandb: Starting sweep agent...
# Runs 5 training jobs with random hyperparameters

Troubleshooting

If wandb.agent does not start, ensure you are logged in with wandb login and your API key is valid.
If the sweep does not optimize as expected, verify the metric name matches exactly and the goal is set correctly to minimize or maximize.
For network issues, check firewall or proxy settings blocking wandb communication.

✅

Key Takeaways

Use method: bayes in sweep config to enable Bayesian optimization with wandb sweeps.
Define clear metric and parameter ranges to guide the optimization effectively.
Run wandb.agent with your training function to execute the sweep and log results automatically.

Verified 2026-04

Verify ↗