> ## Documentation Index
> Fetch the complete documentation index at: https://wb-21fd5541-sweeps-updates.mintlify.site/llms.txt
> Use this file to discover all available pages before exploring further.

> Add W&B to your Python code script or Jupyter Notebook.

# 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.

<Note>
  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.
</Note>

```python lines title="train.py" theme={null}
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](/models/ref/python) to start, stop, and manage sweeps, follow the instructions in the **Python script or notebook** tab. To use the [W\&B CLI](/models/ref/cli) instead, follow the instructions in the **CLI** tab.

<Tabs>
  <Tab title="CLI">
    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.

    <Info>
      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`.
    </Info>

    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`.

    ```yaml lines title="config.yaml" theme={null}
    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](/models/sweeps/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:

    1. Import the W\&B Python SDK (`wandb`).
    2. Initialize a run with [`wandb.init()`](/models/ref/python/functions/init).
    3. Access the hyperparameter values from [`wandb.Run.config`](/models/ref/python/experiments/run#param-config) so that your script uses the suggested arguments for each run instead of hard-coded values.
    4. Log the metric you want to optimize with [`wandb.Run.log()`](/models/ref/python/experiments/run#method-run-log).

    <Important>
      You must log the metric you defined in your configuration.
    </Important>

    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.

    ```python lines title="train.py" theme={null}
    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()
    ```

    <Note>
      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`](/models/ref/cli/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.
    </Note>

    1. Initialize the sweep with the [`wandb sweep`](/models/ref/cli/wandb-sweep) command. Provide the name of the YAML file. Optionally, set the `--project` flag to the project name:

       ```bash theme={null}
       wandb sweep --project project_name config.yaml
       ```

    2. 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`](/models/ref/cli/wandb-agent) to start the sweep agent:

       ```bash theme={null}
       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`:

       ```bash theme={null}
       wandb agent --count 4 entity_name/project_name/sweep_id
       ```

    For more information, see [Start a sweep agent](/models/sweeps/start-sweep-agents).
  </Tab>

  <Tab title="Python script or notebook">
    Follow these steps to add W\&B to your Python script:

    1. Create a dictionary object where the key-value pairs define a [sweep configuration](/models/sweeps/define-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`.
    2. Pass the sweep configuration dictionary to [`wandb.sweep()`](/models/ref/python/functions/sweep). This initializes the sweep and returns a sweep ID (`sweep_id`). For more information, see [Initialize sweeps](/models/sweeps/initialize-sweeps).
    3. At the top of your script, import the W\&B Python SDK (`wandb`).
    4. Within your `main` function, use [`wandb.init()`](/models/ref/python/functions/init) to generate a background process to sync and log data as a [W\&B Run](/models/ref/python/experiments/run). The sweep agent automatically passes the suggested arguments for each run to `wandb.Run.config`.
    5. Fetch the hyperparameter values from `wandb.Run.config`. This lets you use the suggested arguments for each run instead of hardcoded values.
    6. 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.
    7. Start the sweep with [`wandb.agent()`](/models/ref/python/functions/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:

    ```python theme={null}
    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.
  </Tab>
</Tabs>

<Note>
  **Logging metrics to W\&B in a sweep**

  You 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.

  ```python theme={null}
  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.

  ```python theme={null}
  with wandb.init() as run:
      val_loss, val_acc = train()
      run.log(
          {
              "validation": {
                  "val_loss": val_loss, 
                  "val_acc": val_acc,
              }
          }
      )
  ```
</Note>