Add W&B (wandb) to your code

This guide shows how to add W&B Sweeps to an existing Python training script or notebook for hyperparameter optimization. You’ll start with an example training script, then adapt it to explore hyperparameter values, log metrics, and identify the best-performing configuration.

Original training script

Suppose you have a Python script that trains a model (see the following code). Your goal is to find the hyperparameters that maximize the validation accuracy (val_acc). In your Python script, you define two functions: train_one_epoch() and evaluate_one_epoch(). The train_one_epoch() function simulates training for one epoch and returns the training accuracy and loss. The evaluate_one_epoch() function simulates evaluation of the model on the validation data set and returns the validation accuracy and loss. You define a configuration dictionary named config that contains hyperparameter values such as the learning rate, batch size, and number of epochs. The values in the configuration dictionary control the training process. Next, you define a function called main that mimics a typical training loop. For each epoch, the script computes the accuracy and loss on the training and validation data sets.

This code is a mock training script. It doesn’t train a model, but simulates the training process by generating random accuracy and loss values. The purpose of this code is to demonstrate how to integrate W&B into your training script.

train.py

import random
import numpy as np

def train_one_epoch(epoch, lr, batch_size):
    acc = 0.25 + ((epoch / 30) + (random.random() / 10))
    loss = 0.2 + (1 - ((epoch - 1) / 10 + random.random() / 5))
    return acc, loss

def evaluate_one_epoch(epoch):
    acc = 0.1 + ((epoch / 20) + (random.random() / 10))
    loss = 0.25 + (1 - ((epoch - 1) / 10 + random.random() / 6))
    return acc, loss

# config variable with hyperparameter values
config = {"lr": 0.0001, "batch_size": 16, "epochs": 5}

def main():
    lr = config["lr"]
    batch_size = config["batch_size"]
    epochs = config["epochs"]

    for epoch in np.arange(1, epochs):
        train_acc, train_loss = train_one_epoch(epoch, lr, batch_size)
        val_acc, val_loss = evaluate_one_epoch(epoch)

        print("epoch: ", epoch)
        print("training accuracy:", train_acc, "training loss:", train_loss)
        print("validation accuracy:", val_acc, "validation loss:", val_loss)

if __name__ == "__main__":
    main()

The following section shows how to add W&B to your Python script to track hyperparameters and metrics during training. The goal is to find the best hyperparameters that maximize the validation accuracy (val_acc).

Add W&B to your training script

This section shows how to modify the original training script so that the sweep agent can pass hyperparameter values into each run and W&B can record the resulting metrics. How you integrate W&B into your Python script or notebook depends on how you manage sweeps. To use the W&B Python SDK to start, stop, and manage sweeps, follow the instructions in the Python script or notebook tab. To use the W&B CLI instead, follow the instructions in the CLI tab.

CLI
Python script or notebook

Create a YAML configuration file with your sweep configuration. W&B uses this file to determine which hyperparameters and metric to optimize.Add the name of your Python script to the program key in the YAML file on line 1.

The sweep agent selects a value from the values list and passes it to the run config in the training script. For example, if you define the batch_size parameter with the values [16, 32, 64], the sweep agent selects one of those values and passes it to the training script as run.config.batch_size.

The following YAML file replicates to the config values in the Python script (see line 15) shown earlier. The YAML file defines the batch_size, lr, and epochs hyperparameters and specifies the values to try for each one on lines 8–14. On line 5, the YAML file configures the sweep to maximize val_acc.

config.yaml

program: train.py
method: random
name: sweep
metric:
  goal: maximize
  name: val_acc
parameters:
  batch_size:
    values: [16, 32, 64]
  lr:
    min: 0.0001
    max: 0.1
  epochs:
    values: [5, 10, 15]

For more information, see Define sweep configuration.After you define your sweep configuration in a YAML file, add W&B to your training script so that each sweep run can use the hyperparameters selected by the sweep agent and log the metric you want to optimize.Within your training script, add the following code snippets to integrate W&B:

Import the W&B Python SDK (wandb).
Initialize a run with wandb.init().
Access the hyperparameter values from wandb.Run.config so that your script uses the suggested arguments for each run instead of hard-coded values.
Log the metric you want to optimize with wandb.Run.log().

The following code snippet shows how to integrate W&B into your training script. When the sweep agent runs this script, it passes the selected hyperparameter values to wandb.Run.config for that run.

train.py

import wandb
import random
import numpy as np

def train_one_epoch(epoch, lr, batch_size):
    """Simulates training for one epoch and returns the training accuracy and loss."""
    acc = 0.25 + ((epoch / 30) + (random.random() / 10))
    loss = 0.2 + (1 - ((epoch - 1) / 10 + random.random() / 5))
    return acc, loss

def evaluate_one_epoch(epoch):
    """Simulates evaluation for one epoch and returns the validation accuracy and loss."""
    acc = 0.1 + ((epoch / 20) + (random.random() / 10))
    loss = 0.25 + (1 - ((epoch - 1) / 10 + random.random() / 6))
    return acc, loss

def main():
    with wandb.init() as run:
        lr = run.config["lr"]
        batch_size = run.config["batch_size"]
        epochs = run.config["epochs"]

        for epoch in np.arange(1, epochs):
            train_acc, train_loss = train_one_epoch(epoch, lr, batch_size)
            val_acc, val_loss = evaluate_one_epoch(epoch)
            run.log(
                {
                    "epoch": epoch,
                    "train_acc": train_acc,
                    "train_loss": train_loss,
                    "val_acc": val_acc,
                    "val_loss": val_loss,
                }
            )

# Call the main function.
main()

When you create and manage a sweep with the W&B CLI, do not read the sweep configuration file from your training script. Pass the YAML file to wandb sweep when you create the sweep. When a sweep agent starts a run, it automatically populates wandb.Run.config with the selected hyperparameter values.If you run the training script directly with python train.py, no sweep agent is present to populate those values. As a result, keys such as wandb.Run.config["lr"] are unavailable.

Initialize the sweep with the wandb sweep command. Provide the name of the YAML file. Optionally, set the --project flag to the project name:
wandb sweep --project project_name config.yaml
Copy the sweep ID. Replace the placeholder values (entity_name, project_name, sweep_id) in the following command with your W&B entity, project, and sweep ID, then run wandb agent to start the sweep agent:
wandb agent entity_name/project_name/sweep_id
If you want to limit how many sweep runs the agent executes, specify an integer with --count:
wandb agent --count 4 entity_name/project_name/sweep_id

For more information, see Start a sweep agent.

Follow these steps to add W&B to your Python script:

Create a dictionary object where the key-value pairs define a sweep configuration. The sweep configuration defines the hyperparameters you want W&B to explore along with the metric you want to optimize. Continuing from the previous example, vary the batch_size, epochs, and lr hyperparameters during each sweep. To maximize validation accuracy, set the metric’s goal to maximize and its name to val_acc.
Pass the sweep configuration dictionary to wandb.sweep(). This initializes the sweep and returns a sweep ID (sweep_id). For more information, see Initialize sweeps.
At the top of your script, import the W&B Python SDK (wandb).
Within your main function, use wandb.init() to generate a background process to sync and log data as a W&B Run. The sweep agent automatically passes the suggested arguments for each run to wandb.Run.config.
Fetch the hyperparameter values from wandb.Run.config. This lets you use the suggested arguments for each run instead of hardcoded values.
Log the metric you’re optimizing for to W&B using wandb.Run.log(). You must log the metric defined in your configuration. For example, if you define the metric to optimize as val_acc, you must log val_acc. If you don’t log the metric, W&B can’t perform optimization. Within the configuration dictionary (sweep_configuration in this example), you define the sweep to maximize the val_acc value.
Start the sweep with wandb.agent(). Provide the sweep ID and the name of the function the sweep executes (function=main), and set the maximum number of runs to four (count=4).

Your script might look similar to the following:

import wandb # Import the W&B Python SDK
import numpy as np
import random

def train_one_epoch(epoch, lr, batch_size):
    acc = 0.25 + ((epoch / 30) + (random.random() / 10))
    loss = 0.2 + (1 - ((epoch - 1) / 10 + random.random() / 5))
    return acc, loss

def evaluate_one_epoch(epoch):
    acc = 0.1 + ((epoch / 20) + (random.random() / 10))
    loss = 0.25 + (1 - ((epoch - 1) / 10 + random.random() / 6))
    return acc, loss

def main():
    with wandb.init() as run:
        # Fetch the hyperparameter values from run.config
        lr = run.config["lr"]
        batch_size = run.config["batch_size"]
        epochs = run.config["epochs"]

        # Execute the training loop and log the performance values to W&B
        for epoch in np.arange(1, epochs):
            train_acc, train_loss = train_one_epoch(epoch, lr, batch_size)
            val_acc, val_loss = evaluate_one_epoch(epoch)
            run.log(
                {
                    "epoch": epoch,
                    "train_acc": train_acc,
                    "train_loss": train_loss,
                    "val_acc": val_acc, # Metric optimized
                    "val_loss": val_loss,
                }
            )

if __name__ == "__main__":
    # Define a sweep config dictionary
    sweep_configuration = {
        "method": "random",
        "name": "sweep",
        # Metric that you want to optimize
        # For example, if you want to maximize validation
        # accuracy set "goal": "maximize" and the name of the variable 
        # you want to optimize for, in this case "val_acc"
        "metric": {
            "goal": "maximize",
            "name": "val_acc"
            },
        "parameters": {
            "batch_size": {"values": [16, 32, 64]},
            "epochs": {"values": [5, 10, 15]},
            "lr": {"max": 0.1, "min": 0.0001},
        },
    }

    # Initialize the sweep by passing in the config dictionary
    sweep_id = wandb.sweep(sweep=sweep_configuration, project="sweep-example")

    # Start the sweep job
    wandb.agent(sweep_id, function=main, count=4)

When you run this script, W&B starts the sweep, calls the main function up to four times with different hyperparameter combinations, and logs each run’s metrics so you can compare results in the W&B App.

Logging metrics to W&B in a sweepYou must name the metric you want to optimize in your sweep configuration and log that same metric key with wandb.Run.log(). For example, if you define the metric to optimize as val_acc within your sweep configuration, you must also log val_acc to W&B. If you don’t log the metric, W&B can’t perform optimization.

with wandb.init() as run:
    val_loss, val_acc = train()
    run.log(
        {
            "val_loss": val_loss,
            "val_acc": val_acc,
        }
    )

The following is an incorrect example of logging the metric to W&B. The sweep configuration optimizes for val_acc, but the code logs val_acc within a nested dictionary under the key validation. You must log the metric directly, not within a nested dictionary.

with wandb.init() as run:
    val_loss, val_acc = train()
    run.log(
        {
            "validation": {
                "val_loss": val_loss, 
                "val_acc": val_acc,
            }
        }
    )

Guides

Integrations

Reference

Original training script

Add W&B to your training script

​Original training script

​Add W&B to your training script

Original training script

Add W&B to your training script