Documentation
English
Search…
⌃K

Add W&B to your code

There are numerous ways to add the Weights & Biases Python SDK to your script or Jupyter Notebook. Outlined below is a "best practice" example of how to integrate the W&B Python SDK into your own code.

Original training script

Suppose you have the following code in a Jupyter Notebook cell or Python script. We define a function called main that mimics a typical training loop; for each epoch, the accuracy and loss is computed on the training and validation data sets. The values are randomly generated for the purpose of this example.
We defined a dictionary called config where we store hyperparameters values (line 15). At the end of the cell, we call the main function to execute the mock training code.
1
#train.py
2
import random
3
import numpy as np
4
5
def train_one_epoch(epoch, lr, bs):
6
acc = 0.25 + ((epoch/30) + (random.random()/10))
7
loss = 0.2 + (1 - ((epoch-1)/10 + random.random()/5))
8
return acc, loss
9
10
def evaluate_one_epoch(epoch):
11
acc = 0.1 + ((epoch/20) + (random.random()/10))
12
loss = 0.25 + (1 - ((epoch-1)/10 + random.random()/6))
13
return acc, loss
14
15
config = {
16
'lr' : 0.0001,
17
'bs' : 16,
18
'epochs': 5
19
}
20
21
def main():
22
# Note that we define values from `wandb.config` instead of
23
# defining hard values
24
lr = config['lr']
25
bs = config['bs']
26
epochs = config['epochs']
27
28
for epoch in np.arange(1, epochs):
29
train_acc, train_loss = train_one_epoch(epoch, lr, bs)
30
val_acc, val_loss = evaluate_one_epoch(epoch)
31
32
print('epoch: ', epoch)
33
print('training accuracy:', train_acc,'training loss:', train_loss)
34
print('validation accuracy:', val_acc,'training loss:', val_loss)
35
36
# Call the main function.
37
main()

Training script with W&B Python SDK

The following code examples demonstrate how to add the W&B Python SDK into your code. If you start W&B Sweep jobs in the CLI, you will want to explore the CLI tab. If you start W&B Sweep jobs within a Jupyter notebook or Python script, explore the Python SDK tab.
Python script or Jupyter Notebook
CLI
To create a W&B Sweep, we added the following to the code example:
  1. 1.
    Line 1: Import the Wights & Biases Python SDK.
  2. 2.
    Line 6: Create a dictionary object where the key-value pairs define the sweep configuration. In the proceeding example, the batch size (batch_size), epochs (epochs), and the learning rate (lr) hyperparameters are varied during each sweep. For more information on how to create a sweep configuration, see Define sweep configuration.
  3. 3.
    Line 19: Pass the sweep configuration dictionary to wandb.sweep. This initializes the sweep. This returns a sweep ID (sweep_id). For more information on how to initialize sweeps, see Initialize sweeps.
  4. 4.
    Line 33: Use the wandb.init() API to generate a background process to sync and log data as a W&B Run.
  5. 5.
    Line 37-39: (Optional) define values from wandb.config instead of defining hard coded values.
  6. 6.
    Line 45: Log the metric we want to optimize with wandb.log. You must log the metric defined in your configuration. Within the configuration dictionary (sweep_configuration in this example) we defined the sweep to maximize the val_acc value).
  7. 7.
    Line 54: Start the sweep with the wandb.agent API call. Provide the sweep ID (line 19), the name of the function the sweep will execute (function=main), and set the maximum number of runs to try to four (count=4). For more information on how to start W&B Sweep, see Start sweep agents.
1
import wandb
2
import numpy as np
3
import random
4
5
# Define sweep config
6
sweep_configuration = {
7
'method': 'random',
8
'name': 'sweep',
9
'metric': {'goal': 'maximize', 'name': 'val_acc'},
10
'parameters':
11
{
12
'batch_size': {'values': [16, 32, 64]},
13
'epochs': {'values': [5, 10, 15]},
14
'lr': {'max': 0.1, 'min': 0.0001}
15
}
16
}
17
18
# Initialize sweep by passing in config. (Optional) Provide a name of the project.
19
sweep_id = wandb.sweep(sweep=sweep_configuration, project='my-first-sweep')
20
21
# Define training function that takes in hyperparameter values from `wandb.config` and uses them to train a model and return metric
22
def train_one_epoch(epoch, lr, bs):
23
acc = 0.25 + ((epoch/30) + (random.random()/10))
24
loss = 0.2 + (1 - ((epoch-1)/10 + random.random()/5))
25
return acc, loss
26
27
def evaluate_one_epoch(epoch):
28
acc = 0.1 + ((epoch/20) + (random.random()/10))
29
loss = 0.25 + (1 - ((epoch-1)/10 + random.random()/6))
30
return acc, loss
31
32
def main():
33
run = wandb.init()
34
35
# note that we define values from `wandb.config` instead
36
# of defining hard values
37
lr = wandb.config.lr
38
bs = wandb.config.batch_size
39
epochs = wandb.config.epochs
40
41
for epoch in np.arange(1, epochs):
42
train_acc, train_loss = train_one_epoch(epoch, lr, bs)
43
val_acc, val_loss = evaluate_one_epoch(epoch)
44
45
wandb.log({
46
'epoch': epoch,
47
'train_acc': train_acc,
48
'train_loss': train_loss,
49
'val_acc': val_acc,
50
'val_loss': val_loss
51
})
52
53
# Start sweep job.
54
wandb.agent(sweep_id, function=main, count=4)
To create a W&B Sweep, we first create a YAML configuration file. The configuration file contains he hyperparameters we want the sweep to explore. In the proceeding example, the batch size (batch_size), epochs (epochs), and the learning rate (lr) hyperparameters are varied during each sweep.
# 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 on how to create a W&B Sweep configuration, see Define sweep configuration.
Note that you must provide the name of your Python script for the program key in your YAML file.
Next, we add the following to the code example:
  1. 1.
    Line 1-2: Import the Wights & Biases Python SDK (wandb) and PyYAML (yaml). PyYAML is used to read in our YAML configuration file.
  2. 2.
    Line 18: Read in the configuration file.
  3. 3.
    Line 21: Use the wandb.init() API to generate a background process to sync and log data as a W&B Run. We pass the config object to the config parameter.
  4. 4.
    Line 25 - 27: Define hyperparameter values from wandb.config instead of using hard coded values.
  5. 5.
    Line 33-39: Log the metric we want to optimize with wandb.log. You must log the metric defined in your configuration. Within the configuration dictionary (sweep_configuration in this example) we defined the sweep to maximize the val_acc value.
1
import wandb
2
import yaml
3
import random
4
import numpy as np
5
6
def train_one_epoch(epoch, lr, bs):
7
acc = 0.25 + ((epoch/30) + (random.random()/10))
8
loss = 0.2 + (1 - ((epoch-1)/10 + random.random()/5))
9
return acc, loss
10
11
def evaluate_one_epoch(epoch):
12
acc = 0.1 + ((epoch/20) + (random.random()/10))
13
loss = 0.25 + (1 - ((epoch-1)/10 + random.random()/6))
14
return acc, loss
15
16
def main():
17
# Set up your default hyperparameters
18
with open('./config.yaml') as file:
19
config = yaml.load(file, Loader=yaml.FullLoader)
20
21
run = wandb.init(config=config)
22
23
# Note that we define values from `wandb.config` instead of
24
# defining hard values
25
lr = wandb.config.lr
26
bs = wandb.config.batch_size
27
epochs = wandb.config.epochs
28
29
for epoch in np.arange(1, epochs):
30
train_acc, train_loss = train_one_epoch(epoch, lr, bs)
31
val_acc, val_loss = evaluate_one_epoch(epoch)
32
33
wandb.log({
34
'epoch': epoch,
35
'train_acc': train_acc,
36
'train_loss': train_loss,
37
'val_acc': val_acc,
38
'val_loss': val_loss
39
})
40
41
# Call the main function.
42
main()
Navigate to your CLI. Within your CLI, set a maximum number of runs the sweep agent should try. This is step optional. In the following example we set the maximum number to five.
NUM=5
Next, initialize the sweep with the wandb sweep command. Provide the name of the YAML file. Optionally provide the name of the project for the project flag (--project):
wandb sweep --project sweep-demo-cli config.yaml
This returns a sweep ID. For more information on how to initialize sweeps, see Initialize sweeps.
Copy the sweep ID and replace sweepID in the proceeding code snippet to start the sweep job with the wandb agent command:
wandb agent --count $NUM noahluna/sweep-demo-cli/sweepID
For more information on how to start sweep jobs, see Start sweep jobs.