Sweep from Jupyter Notebook

Run sweeps from Jupyter notebooks

Initialize a sweep

import wandb
sweep_config = {
"name": "My Sweep",
"method": "grid",
"parameters": {
"param1": {
"values": [1, 2, 3]
sweep_id = wandb.sweep(sweep_config)

Use the following methods in order to specify the entity or project for the sweep:

Run an agent

When running an agent from python, the agent runs a specified function instead of using the program key from the sweep configuration file.

import wandb
import time
def train():
run = wandb.init()
print("config:", dict(run.config))
for epoch in range(35):
print("running", epoch)
wandb.log({"metric": run.config.param1, "epoch": epoch})
wandb.agent(sweep_id, function=train)


Using wandb.agent() with jupyter notebook environments can hang when using GPUs.

There can be a bad interaction between wandb.agent() and jupyter environments due to how GPU/CUDA resources are initialized by frameworks.

A temporary workaround (until we can fix these interactions) is to avoid using the python interface for running the agent. Instead, use the command line interface by setting the program key in the sweep configuration, and execute: !wandb agent SWEEP_ID in your notebook.


  • sweep_id (dict): Sweep ID generated by the UI, CLI, or sweep API

  • entity (str, optional): username or team where you want to send runs

  • project (str, optional): project where you want to send runs

  • function (dir, optional): Configure sweep function

Run a local controller

If you want to develop your own parameter search algorithms you can run your controller from python.

The simplest way to run a controller:

sweep = wandb.controller(sweep_id)

If you want more control of the controller loop:

import wandb
sweep = wandb.controller(sweep_id)
while not sweep.done():

Or even more control over the parameters being served:

import wandb
sweep = wandb.controller(sweep_id)
while not sweep.done():
params = sweep.search()

If you want to specify your sweep entirely with code you can do something like this:

import wandb
from wandb.sweeps import GridSearch,RandomSearch,BayesianSearch
sweep = wandb.controller()
sweep.configure_parameter('param1', value=3)