Sweeps

"""
Example sweep configuration for W&B Sweeps.

sweep_configuration is an example of a single nested sweep configuration
sweep_configuration_nested is an example of a double nested sweep configuration

Replace:
- values enclosed in angle brackets with your own
- integer values, float values, and lists of values with your own hyperparameter search space
"""

sweep_configuration = {
    "program": "<program_name>",
    "name": "<sweep_name>",
    "method": "<sweep_method>",
    "metric": {
        "name": "<metric_name>",
        "goal": "<goal_type>"
        },
    "parameters": {
        "hyperparameter_1": {
            "search_constraint_1": 0.0001, 
            "search_constraint_2": 0.1
            },
        "hyperparameter_2": {
            "values": [16, 32, 64]
            },
        "hyperparameter_3": {
            "values": [5, 10, 15]
            },
        "hyperparameter_4": {
            "values": ["value_1", "value_2"]
            },
    },
}

## Double nested sweep configuration example
sweep_configuration_nested = {
  "program": "<program_name>",
  "name": "<sweep_name>",
  "method": "<sweep_method>",
  "metric": {
    "name": "<metric_name>",
    "goal": "<goal_type>"
  },
  "parameters": {
    "hyperparameter_1": {
      "values": ["<value_1>", "<value_2>"]
    },
    "hyperparameter_2": {
      "values": [128, 256, 512]
    },
    "hyperparameter_3": {
      "values": [0.3, 0.4, 0.5]
    },
    "hyperparameter_4": {
      "value": 1
    },
    "hyperparameter_5": {
      "distribution": "<distribution_type>",
      "search_constraint_1": 0,
      "search_constraint_2": 0.1
    },
    "nested_category_1": {
      "parameters": {
        "nested_hyperparameter_1": {
          "distribution": "<distribution_type>",
          "search_constraint_1": 0.0,
          "search_constraint_2": 0.9
        },
        "nested_hyperparameter_2": {
          "values": [0.0001, 0.0005, 0.001]
        }
      }
    },
    "nested_category_2": {
      "parameters": {
        "nested_hyperparameter_1": {
          "distribution": "<distribution_type>",
          "search_constraint_1": 0.0,
          "search_constraint_2": 0.9
        },
        "nested_hyperparameter_2": {
          "values": [0.1, 0.2, 0.3]
        }
      }
    },
    "nested_category_3": {
      "parameters": {
        "nested_hyperparameter_1": {
          "distribution": "<distribution_type>",
          "search_constraint_1": 0.5,
          "search_constraint_2": 0.7
        },
        "nested_hyperparameter_2": {
          "values": [0.2, 0.3, 0.4]
        }
      }
    }
  }
}

"""
Define a sweep config, initialize, and start W&B Sweep.

sweep_configuration is a single nested sweep configuration. See sweep_config.py
for an example of a double nested sweep configuration.
"""
import wandb
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 will be None,
    # so we 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": {
            "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=2)

"""
Initialize a W&B Sweep.

Replace:
- values enclosed in angle brackets with your own
- sweep_configuration with your own sweep configuration dictionary. See
    sweep_config.py for examples of single nested and double nested sweep
    configurations.
"""

import wandb

# Example sweep configuration
sweep_configuration = {
    "method": "<sweep_method>",
    "name": "<sweep_name>",
    "metric": {
        "goal": "<goal>", 
        "name": "<metric_name>"
        },
    "parameters": {
        "hyperparameter_1": {
            "values": [16, 32, 64]
            },
        "hyperparameter_2": {
            "values": [5, 10, 15]
            },
        "hyperparameter_3": {
            "search_constraint_1": 0.1,
            "search_constraint_2": 0.0001
            },
    },
}

sweep_id = wandb.sweep(sweep=sweep_configuration, project="<project>")

"""
Start a sweep job to run a sweep agent.

Replace:
- values enclosed in angle brackets with your own
- sweep_id with the ID of your sweep, which is returned when you initialize a sweep.
    See sweep_initialize.py for an example of how to initialize a sweep and get
    the sweep_id.
- count with the number of runs you want to execute in your sweep
"""

import wandb

sweep_id, count = "<sweep_id>", 10
wandb.agent(sweep_id, count=count)