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.
What is experiment tracking?
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, …)
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.
How to do 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