Introduction

I often find myself coding a machine learning experiment in a Jupyter Notebook, and at the same time, using Weights & Biases (wandb) to visualize and track the results of the runs. When the experiment is finished, I always have questions such as: How will the performance be affected by the parameter a? What if I change the number of items of the dataset, or change the dataset completely?

Hyperpameter tuning with wandb sweeps is a great tool to solve these questions. However, sweeping requires that you define a specific training program, as a separate python file. I find this to be redundant, specially when the code for training is already in the Jupyter Notebook. Furthermore, if I make some changes in the original notebook, I have to be sure that I change the sweep script too.

This post shows a trick to execute a Jupyter Notebook as the program of a wandb sweep. This provides a frictionless way of using your Jupyter Notebooks both for single runs and sweep functions. We won't use any separate configuration or script file, everything will be done between Jupyter and wandb. This post assumes that the reader has basic knowledge on both how Jupyter Notebooks and wandb sweeps work.

As use case we will perform a time series classification task with deep neural networks using the wonderful library tsai. This is all the code needed to train a classifier in tsai for the dataset NATOPS:

from tsai.all import *

dsid = 'NATOPS' 
X, y, splits = get_UCR_data(dsid, return_split=False)
learn = TSClassifier(X, y, splits=splits, bs=[64, 128], 
                     batch_tfms=[TSStandardize()], arch=InceptionTime, 
                     metrics=accuracy)
learn.fit_one_cycle(25, lr_max=1e-3)
learn.plot_metrics()

We will do a hyperparameter search over the two arguments of the call to fit_one_cycle, that is, the number of epocs (n_epochs) and the maximum learning rate passed to the one-cycle schedule (lr_max).

In the next section, we'll see how to organize the notebook so that it is ready to be used as the program of a sweep. Then, we'll configure it to be run in a local server (e.g, an instance of JupyterLab). Finally, for Colab users, we'll see a workaround to make this work with a subtle difference.

Refactoring the notebook for wandb sweeps

As explained in the wandb documentation: "two components work together in a sweep: a controller on the central sweep server, which picks out new hyperparameter combinations to try, and agents, running in any number of processes on any number of machines, which query the server for hyperparameters, use them to run model training, and then report the results back to the controller."

Each time the agent queries values of the hyperparameters for a new trial, those will be injected as part of the configuration of the wandb run that the training program must have. Once the program executes the call to wandb.init to begin the syncing, the object wandb.config will contain them, and any line of code that depends on that config will use the values pof that trial.

But, what happens if we had already defined a configuration object like the one below, before the call to wandb.init to play manually with different values?

config = {
    'n_epochs' : 25,
    'lr_max' : 1e-3,
    'bs' : 64
}

import wandb
run = wandb.init(config=config, mode='disabled')

The good thing is that, nothing happens! Even if we have set values for n_epochs and lr_max before calling wandb.init, the sweep agent will override them with the values of the new trial. This is a common question explained in the wandb docs, and it is exactly what allows us to use the same exact notebook for both single runs and sweeps. Parameters that are not part of the sweep, such as bs in this example, can be part of the config object as well and of course they will be kept there by the sweep agent.

The last thing we have to do is, as in every sweep, replace our magic numbers (at least the ones that are part of the sweep) with the corresponding reference to the config variable:

from tsai.all import *

dsid = 'NATOPS' 
X, y, splits = get_UCR_data(dsid, return_split=False)
learn = TSClassifier(X, y, splits=splits, bs=[config['bs'], 128], 
                     batch_tfms=[TSStandardize()], arch=InceptionTime, 
                     metrics=accuracy, cbs=[WandbCallback(log_preds=False)])
learn.fit_one_cycle(config['n_epochs'], lr_max=config['lr_max'])
learn.plot_metrics()

So that's basically all you have to do to make your notebook ready for both single experiments & sweepin: move your magic numbers that you want to sweep over to an initial config object, and pass that as config to wandb.init.

Configuring the sweep

There are many ways to create the configuration of a new sweep for wandb:

  • Use the graphical user interface
  • Create a separate yaml file
  • Define it somewhere in your notebook as a dictionary or a JSON object

I like to use directly the graphical interface. In this way, I don't have to create a separate yaml file, and I don't have to touch my notebook, which makes everything cleaner. If you have never created a sweep using the wandb interface, there's a big button "Create new sweep" on the top-right corner of the sweeps tab

You will see a nice YAML editor in which you have to define the parameters of the sweep, as well metric to optimize. You can also assign a name, a description, a method (we will use bayes here) and many more. Below you can see how a sweep to search over the parameters n_epoch and lr_max with respect to the validation loss would look like, but wait until you press the "Initialize sweep" button...

The trick comes with the program attribute. By default, wandb expects that you have defined a Python script called train.py, that contains your experiment synchronized with wandb. Our program is a Jupyter notebook though, so we will change this with the absolute path of our notebook in our Jupyter server.

program: /home/victor/work/_notebooks/2021-09-26-sweeps.ipynb

The question is: How are we going to execute the notebook? Obviously, if we run the sweep like this, wandb will try to execute the notebook as a Python script and the agent will crash. To solve this, we use the command key. This configuration key tells the wandb agent the command structure for invoking and passing arguments to the training script. By default, it is defined as:

command:
  - ${env}
  - ${interpreter}
  - ${program}
  - ${args}

where ${env} is /usr/bin/env (in UNIX systems), ${interpreter} expands to python, ${program} is the file path of our training script (a notebook in our case), and ${args} contain possible parameters of the classic form --param1=value1.

However, we can redefine the command as we wish. More specifically, to execute a notebook, we can make use of the nbconvert tool, part of the Jupyter ecosystem. The exact shell command that we have to type to execute the notebook is:

jupyter nbconvert --to notebook --execute ${program}

Since the sweep expects the command to be in exec form instead of shell form, we will add it to the sweep config as:

command: ["jupyter", "nbconvert", "--to", "notebook", "--execute", "${program}"]

And that's it! Now you can press the big blue "Initialize sweep button", and wandb will prompt you with a command to start an wandb agent that runs the sweep:

Just copy that wandb agent into a terminal in your server and the sweep will start. You can create multiple instances of your agent in one or multiple machines. I use this a lot to use at once all the GPUs of my system in the sweep. As explained in this blog post, it is just a matter of fixing the env variable CUDA_VISIBLE_DEVICES in each of the calls to the agent:

$ CUDA_VISIBLE_DEVICES=0 wandb agent vrodriguezf90/dummy_sweep/gs9p78yg

$ CUDA_VISIBLE_DEVICES=1 wandb agent vrodriguezf90/dummy_sweep/gs9p78yg

Conclusions

If you are like me, and always procastinate doing hyperparameter tuning because of the extra boilerplate needed to make it work, this can be really useful when you work in Jupyter Notebooks. Since this post is in itself a Jupyter Notebook (Yes, you can write blog posts using Jupyter with this awesome tool!), I used it as a program for the sweep that was described in the previous section, and everything worked like a charm. You can visualize the sweep here.

Finally, it is worth to mention that the use of the tool nbconvert in each trial of the sweep creates a bit of overhead, which can be annoying, specially for small sweeps. There are multiple options to overcome this overhead, such as transforming the notebook into a script before configuring the sweep, or using faster tools to convert the notebook into a script such as the function nb2py from the tsai library.

Acknowledgments

Thakn you to Ignacio Oguiza (El Gurú), for encouraging me to write this blog post, and for all the wonderful work and knowledge he puts into the tsai library.