> ## Documentation Index
> Fetch the complete documentation index at: https://docs.wandb.ai/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 provides recommendations on how to integrate W\&B into your Python training script or notebook for hyperparameter search optimization. By following these recommendations, you can use W\&B Sweeps to explore hyperparameter values, log training and validation metrics, and identify the configuration that produces strong model performance.

This guide is for machine learning practitioners who already have a Python training script and want to add hyperparameter sweep support. The following sections walk through an example training script, then show how to update it to work with W\&B Sweeps.

## 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 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. You want to use W\&B 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.

<Tabs>
  <Tab title="CLI">
    Create a YAML configuration file with your sweep configuration. The configuration file contains the hyperparameters you want the sweep to explore. In the following example, the sweep varies the `batch_size`, `epochs`, and `lr` hyperparameters during each run.

    ```yaml theme={null}
    # 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](/models/sweeps/define-sweep-configuration).

    You must provide the name of your Python script for the `program` key
    in your YAML file.

    Next, add the following to the code example:

    1. Import the W\&B Python SDK (`wandb`) and PyYAML (`yaml`). Use PyYAML to read in your YAML configuration file.
    2. Read in the configuration file.
    3. Use [`wandb.init()`](/models/ref/python/functions/init) to start a background process to sync and log data as a [W\&B Run](/models/ref/python/experiments/run). Pass the config object to the config parameter.
    4. Define hyperparameter values from `wandb.Run.config` instead of using hardcoded values.
    5. Log the metric you want to optimize with [`wandb.Run.log()`](/models/ref/python/experiments/run.md/#method-runlog). You must log the metric defined in your configuration. Within the configuration dictionary (`sweep_configuration` in this example), you define the sweep to maximize the `val_acc` value.

    ```python theme={null}
    import wandb
    import yaml
    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


    def main():
        # Set up your default hyperparameters
        with open("./config.yaml") as file:
            config = yaml.load(file, Loader=yaml.FullLoader)

        with wandb.init(config=config) as run:
            for epoch in np.arange(1, run.config['epochs']):
                train_acc, train_loss = train_one_epoch(epoch, run.config['lr'], run.config['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()
    ```

    After you update your training script, initialize and start the sweep from your CLI:

    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 sweep-demo-cli config.yaml
       ```

       This returns a sweep ID. For more information, see [Initialize sweeps](/models/sweeps/initialize-sweeps).

    2. Copy the sweep ID and replace `[SWEEP-ID]` in the following command to start the sweep job with the [`wandb agent`](/models/ref/cli/wandb-agent) command. Replace `[YOUR-ENTITY]` with your W\&B entity name. Optionally, set `--count` to limit the number of runs the agent tries (this example sets it to five):

       ```bash theme={null}
       wandb agent --count 5 [YOUR-ENTITY]/sweep-demo-cli/[SWEEP-ID]
       ```

    The sweep agent runs your training script repeatedly, each time with a different combination of hyperparameter values from your YAML configuration, and logs the results to W\&B. For more information, see [Start sweep jobs](/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). Pass the project name as a parameter to the `wandb.init()` method. If you don't pass a project name, W\&B uses the default project name.
    5. Fetch the hyperparameter values from the `wandb.Run.config` object. This lets you use the hyperparameter values defined in the sweep configuration dictionary instead of hardcoded values.
    6. Log the metric you're optimizing for to W\&B using [`wandb.Run.log()`](/models/ref/python/experiments/run.md/#method-runlog). 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`).

    To put this together, 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
    import argparse

    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(args=None):
        # When called by sweep agent, args is None,
        # so use the project from sweep config
        project = args.project if args else None
        
        with wandb.init(project=project) as run:
            # Fetches the hyperparameter values from `wandb.Run.config` object
            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__":
        parser = argparse.ArgumentParser()
        parser.add_argument("--project", type=str, default="sweep-example", help="W&B project name")
        args = parser.parse_args()

        # 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=args.project)

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

    When you run this script, W\&B starts the sweep, executes 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 log the metric you define and are optimizing for in both your sweep configuration and 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>