Experiment tracking in machine learning | by Dipam Vasani | Towards Data Science

As machine learning matures, we come across newer problems, and then we come up with sophisticated solutions for these problems. For example, in the beginning, it was hard to train a neural network itself. Then came GPUs and frameworks that made this task really easy. After that came Tensorboard for easy metric tracking, Streamlit for creating demo apps, and Docker for easy deploying. The field has come a long way. But there are still many components that need work.
One such component is experiment tracking. We still use physical notebooks or excel sheets to track experiments and then manually plot them. It’s time we changed this.
Training a model is an iterative process. We conduct a bunch of experiments and come up with certain conclusions. These conclusions help us either design further experiments or come to a final inference. Hence, it is essential that we track our experiments well.
Problems with current approaches?
With notebooks or excel sheets,
We don’t have a centralized system to store all our projects and their experiments. They’re also difficult to share or collaborate with others.We often don’t track all parameters and metrics because it is tiresome to do so.We don’t track system metrics like GPU utilization or GPU temperature.We have to write separate code to plot our experiments.Saving checkpoints for each experiment is not easy (model_256.pth, model_256_lr001.pth, ...)
We don’t have a centralized system to store all our projects and their experiments. They’re also difficult to share or collaborate with others.
We often don’t track all parameters and metrics because it is tiresome to do so.
We don’t track system metrics like GPU utilization or GPU temperature.
We have to write separate code to plot our experiments.
Saving checkpoints for each experiment is not easy (model_256.pth, model_256_lr001.pth, ...)
Hence, it would be great if we had a tool for all of this. Fortunately, there is. In this article, I will be using one such tool known as “Weights & Biases”. It is developed by the founders of Spacy and is a really sophisticated framework for ML project management. We will be training a FC neural network for MNIST classification and track it using wandb.
Start by creating a free account on the Weights & Biases website.Open Google colab with the same Google account you used for W&B and create a new notebook.In this article, I will be using ColabCode, a library to use VS Code in Google colab. Start by installing it.
Start by creating a free account on the Weights & Biases website.
Open Google colab with the same Google account you used for W&B and create a new notebook.
In this article, I will be using ColabCode, a library to use VS Code in Google colab. Start by installing it.
pip install colabcode
and then run
Clicking on the ngrok link will give you access to VS Code. You can optionally pass arguments like port or password to ColabCode. Check the Github repository for more details.
4. In your VS Code, open a new terminal and install W&B.
5. Now, on the W&B website, create a new project from the left panel. Here I’ve created a private project called mnist_training.
6. Now, come back to VS code, open a new file, and write your regular train code. I will not go into the details of the PyTorch code but just show the parts related to W&B. You can find the train file here. For this example, I am using only 3 arguments, but you can track all the arguments you want.
7. The first thing we do is import and initialize wandb. This is done using the function wandb.init().
import wandbwandb.init(project = "mnist_training", config = args)
Here config is a dictionary of arguments we want to track. We can either set it during init or update it later. Read the documentation of wandb.config for the various options available.
Next, we use wandb.log to log our metrics.
wandb.log({“train loss”: loss})
This will help you compare the loss across various runs and see how it is affected by other parameters. You can also log images, audios, and specify the type of chart to plot. Again, read the documentation of wandb.log for more information.
8. Saving the checkpoint
To track your models and save your checkpoints, we have two methods.
9. You can run your train script and refresh your W&B page.
It will show you live metric tracking in the Charts section and a bunch of other plots. On the left, you can see all your runs. You can filter them based on date, the ones that did not crash, and so on.
With the table view, you can also assign tags to your experiments for better sorting and grouping.
Finally, you can also perform sweeps to automate hyperparameter optimization. It also has data and model versioning. It just adds so much value to your project.
The best thing about W&B is how easily you can integrate it with any of your existing projects. It supports all modern frameworks for ML and DL. Start using it, and I’m sure you’ll love it.
docs.wandb.com
Another logger you might want to try is called Neptune. As of writing this, many frameworks provide integrations for these loggers. This just goes to show that tracking experiments is a real problem and that the ML community is embracing tools like Neptune and W&B.