W&B Sweeps

sweep_config = {
    "method": "bayes",        # bayes, grid, or 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]
        },
        "epochs": {
            "value": 50             # fixed value
        },
        "dropout": {
            "distribution": "uniform",
            "min": 0.1,
            "max": 0.5
        },
        "optimizer": {
            "values": ["adam", "sgd", "adamw"]
        },
        "hidden_size": {
            "distribution": "int_uniform",
            "min": 64,
            "max": 512
        }
    },
    "early_terminate": {
        "type": "hyperband",
        "min_iter": 5,
        "eta": 3
    }
}

import wandb

def train():
    """Training function called by each sweep agent."""
    run = wandb.init()
    config = wandb.config

    # Build model with sweep parameters
    model = build_model(
        hidden_size=config.hidden_size,
        dropout=config.dropout
    )

    optimizer = get_optimizer(
        config.optimizer,
        model.parameters(),
        lr=config.learning_rate
    )

    for epoch in range(config.epochs):
        train_loss = train_epoch(model, optimizer, config.batch_size)
        val_loss, val_acc = validate(model)

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

    wandb.finish()

# Create the sweep
sweep_id = wandb.sweep(sweep_config, project="sweep-demo")

# Launch agents (run 50 trials)
wandb.agent(sweep_id, function=train, count=50)

# Machine 1: Create the sweep
wandb sweep sweep_config.yaml
# Output: Created sweep with ID: abc123

# Machine 1, 2, 3, ...: Start agents
wandb agent your-entity/your-project/abc123
wandb agent your-entity/your-project/abc123
wandb agent your-entity/your-project/abc123

# Each agent pulls configs from the central sweep controller
# and reports results back. The Bayesian optimizer uses all
# results to choose the next config intelligently.

program: train.py
method: bayes
metric:
  name: val/accuracy
  goal: maximize
parameters:
  learning_rate:
    distribution: log_uniform_values
    min: 0.00001
    max: 0.1
  batch_size:
    values: [16, 32, 64, 128]
  dropout:
    distribution: uniform
    min: 0.0
    max: 0.5
early_terminate:
  type: hyperband
  min_iter: 5