What does Weights & Biases do? MLOps | Omer Mahmood  | Towards Data Science

Weights & Biases (W&B) is a machine learning platform geared towards developers for building better models faster.
It is designed to support and automate key steps in the MLOps life cycle, such as experiment tracking, dataset versioning and model management.
Creating an ML model is the easy part — operationalising and managing the lifecycle of ML models, data and experiments is where things get complicated. That’s why we need a methodical approach and set of practices to address this challenge.
🧐 Still want to know more about MLOps before diving into what W&B does? Here’s a post I wrote previously: What’s MLOps
As with my previous product and platform related posts, I decided to roll up my sleeves one afternoon and give W&B a spin!
W&B is essentially a web-based subscription service, personal accounts are free and come with a 100 GB allowance for dataset storage and artifact tracking. You can sign up here: https://wandb.ai/site
There are two main components that make up W&B, the Workspace and a Python API. In fact, below is a nice summary of the capabilities offered via these two components.
The top row of boxes represents enablers for MLOps, the bottom row shows the corresponding W&B feature, I will touch on them all in this post.
When you sign into W&B, you’ll typically see a Dashboard like the one below, it has a simple clean navigation bar on the left where you can access recent Projects or create a new one. You can also view Teams you are a member of, Resources and Settings.
The main pane on the right of the navigation bar, contains a recent list of ‘Runs’, more on those later.
This is how W&B integrates with your existing model code, and is able to surface the insights from your experiments in the Workspace.
There are quickstart guides depending on your preferred framework, in this post I went with PyTorch. I used one of the PyTorch examples in the W&B github repo for training a CNN model using the Fashion MNIST dataset. This is an MNIST-like dataset, but instead of using it to train a neural network model for classifying numeric digits, the aim is to identify common categories of clothing, accessories and footwear, such as boots, sneakers, coats and bags.
It really is quite straightforward to integrate W&B with a few lines of code. I won’t regurgitate the quickstart here, but in a few steps:
Install the W&B Python API libraries in your environment and login to your W&B account (auth token supplied in the quickstart).
Install the W&B Python API libraries in your environment and login to your W&B account (auth token supplied in the quickstart).
pip install wandbwandb login
2. Add an Import line for the W&B library in your existing model training code, define the project in which you want to track experiment runs, and configure hyperparameters.
import wandbwandb.init(project="medium-post", entity="omerm")wandb.config = {"learning_rate": 0.01,"epochs": 10,...}
3. Finally add a line of code to your training loop to log the metrics you want to track and visualise in your Workspace.
wandb.log({"loss": loss})
One of the key pillars of MLOps is the ability to carry out various experiments, and have some way of tracking how each different configuration affects the model metrics you are most interested in.
When you run the code that contains the training process for a given experiment, the Runs section in a Project is where W&B surfaces all this good experiment data, as shown in the screenshot below.
The Runs view shows you the list of experiment Runs on the left, with chart plots on the right representing the metrics you chose to log from each one.
You have the ability to search for a Run, filter and conveniently group by Team or metric for example.
Metrics that are common across multiple experiment Runs are represented by different coloured plots for each experiment.
Clicking on an individual Run let’s you drill into the further detailed configuration, metrics and various other statistics related to the experiment, as shown in the screenshot below.
Another common MLOps task that falls under model experimentation and development is hyperparameter tuning.
My go-to analogy for hyperparameter tuning is cake baking 🍰. Consider for a moment that the framework, training data and code you write for building a model are the ingredients you combine to make the cake mixture.
You then run that code in order to train a model, or indeed put the cake mix in the oven to bake.
You can think of hyperparameters as the knobs on the oven that control how the cake is baked, for example the timer, temperature, convection type etc. Changing those knobs can have a big impact on how your cake will turn out; too hot and it will burn, if you don’t cook for long enough it might not rise.
The same can be said for models — the hyperparameter values, such as learning rate, and batch size can have a big impact on how long a model takes to converge (or minimize loss).
Loss is the penalty for a bad prediction. That is, loss is a number indicating how bad the model’s prediction was on a single example.[2]
If the model’s prediction is perfect, the loss is zero; otherwise, the loss is greater. The goal of training a model is to find a set of weights and biases 😏 that have low loss, on average, across all examples.[2]
In W&B we can use ‘Sweeps’ to carry out hyperparameter tuning jobs, and easily distribute them across multiple machines to speed up the process.
It’s easy to create and configure a sweep from the Workspace (as shown in figures 4 & 5) or via an API call from your code.
Once you’ve set the parameters of your sweep, you can start it with a single-line command, then monitor its progress and specifically the runs completed so far.
You get some nice visualisations that summarise the results of your sweeps. For example, a chart highlighting parameter importance with respect to particular metrics such as loss or accuracy.
My personal favourite is the parallel coordinates chart, which for example can show the parameter value for each sweep and its impact on the accuracy of a particular label as shown below:
Ultimately you can use the output from your sweeps to choose the optimal hyperparameter values for your model training.
💡 Pro tip, as I discovered while playing around, sweeps for Bayesian or Random parameter searches will run forever unless you force stop (ctrl-c) the process from the command line or Workspace UI. Or you can specify the number of runs an agent should attempt upfront using the count flag:
NUM=10SWEEPID="ex56dcvv"wandb agent --count $NUM $SWEEPID
I’m not going to spend too much time running through this set of functionality as the W&B docs do a pretty good job of walking you through.
Reports are a fundamental part of what makes W&B useful for sharing findings and collaborating across teams.
You get a super rich and flexible Notion-like “Type ‘/’ for commands” UI for creating a report. It is made up of panels that can contain anything from text and markdown to pre-built charts that pull data directly from the various runs in your project.
One of my favourite report elements is the ‘Text table’ panel. It’s a super useful way to visualise data you logged via the wandb.log() function, such as the results from your model evaluation.
Following on with the Fashion MNIST example, I can view the predictions my model made for each image in the dataset. I can also interactively filter the table to show me incorrect predictions, guess vs. truth, as shown below:
It really does take away the frustration from having to curate, summarise and share your data science projects findings quickly, and does it in a beautifully clean way.
The W&B Workspace includes an Artifact registry. Here you can see all artifacts related to a specific project, you can add datasets and models to this view by instantiating a W&B artifact object and saving it in your code, for example:
artifact = wandb.Artifact(name, “model”)artifact.save()
Anything you save is then visible in your project workspace as shown below:
Put simply, the W&B API makes it easy to store metadata about your datasets and models. There is also a useful Graph view that illustrates the lineage of your artifacts.
The ability to track artifacts is useful, but how does this help support our MLOps pipeline when it comes to monitoring and deployment into production? This is where W&B’s model management features come in handy!
At the time of writing, W&B’s Model management features are in beta, and they are ‘actively building out the model registry and model evaluation use cases’.[1]
W&B for Model Management capabilities are designed to track and report on the complete lifecycle of a model:
Datasets: The exact version of the dataset a model trained on
Code: The code used in model training
Models: The weights of the trained model itself
Metrics: The evaluation results of a model on different golden datasets
Statuses: Where each model is in the pipeline (ex. “staging” or “production”)[1]
This list certainly ticks a lot of boxes when it comes to implementing an MLOps pipeline. I gave the beta capabilities a quick spin using the github example in the docs.
The example code basically covers the following steps:
Generate a dataset for model training and evaluation (download of the MNIST data) and publish it to the Artifacts store.Train several models using the dataset from step 1, varying a hyperparameter to differentiate each one. This step also publishes each candidate model and it’s metrics to the Artifacts store in creation order.Evaluate each model to determine its loss; the model with the lowest loss is then ‘promoted’ to production.
Generate a dataset for model training and evaluation (download of the MNIST data) and publish it to the Artifacts store.
Train several models using the dataset from step 1, varying a hyperparameter to differentiate each one. This step also publishes each candidate model and it’s metrics to the Artifacts store in creation order.
Evaluate each model to determine its loss; the model with the lowest loss is then ‘promoted’ to production.
⚠️ Keep in mind that when I say ‘promoted’ there is no actual model deployment happening in your production environment. This is purely a metadata tag that can be applied to an artifact for the purposes of monitoring and tracking — but that’s not to say you couldn’t trigger a script to carry out real actions based on metadata changes.
You can then visualize the results using a Report and a Panel containing a Weave component as shown below:
The weave component is similar to a text table, except that it is used to pull data from the Artifact store. You can then use this as part of a Report, surfacing useful information such as date of creation, loss, status etc.
I never set out to write a comprehensive product review 😅. But I wanted to at least articulate what W&B does and share my brief experience of using it! Here’s a summary...
Creating a free account was very straightforward using a Google account. I had no issue authenticating to W&B from the Google Colab notebook I was using to test.
I didn’t delve into the more advanced features, such as distributed Sweeps or adding multiple contributors to a project.
W&B is definitely not a platform for absolute beginners. You won’t find any pre-built models, or no-code/AutoML options for building models here.
However, if you are in (or manage) a well established team of data scientists/engineers writing python code, and have been looking for a more structured way to collaborate across projects, while also driving an MLOps culture — W&B could be a good fit.
It’s striking how easy it is to start getting value from W&B. With just a few lines of code and broad support for the most popular ML frameworks, you should be able to integrate it with your existing python data science workflow very quickly.
I would also say that the documentation is excellent, with lots of videos, worked examples in Google Colab notebooks and a great user community.
The free, personal account offers a generous 100 GB for project-related storage and artifacts; the compute resource required to make the W&B workspace work also seems to be included.
There are concessions for academic and open source teams. Beyond that, they charge on a ‘pay-as-you-go’ basis depending on how much data you want to store and number of artifacts you want to track.
Just remember you still need an environment for running your code, meaning you’re on the hook for any resources consumed in the process of model training or evaluation.
I really enjoyed testing out W&B for this post. I think as the product evolves it could become a more of an end-to-end platform for automating MLOps processes, but right now it does a pretty good job enabling an MLOps culture.
If you’ve reached the end then thanks for reading and I’ll see you on the next one! 👋
[1] W&B, Documentation, [Beta] Model Management: https://docs.wandb.ai/guides/models
[2] Google Machine Learning Crash Course, Descending into ML: Training and Loss: https://developers.google.com/machine-learning/crash-course/descending-into-ml/training-and-loss