The Best Tools for Machine Learning Model Visualization

The phrase “Every model is wrong, but some are useful” is especially true in Machine Learning. When developing machine learning models, you should always understand where it works as expected and where it fails miserably.

An important method of understanding a model is through model visualization. Visualizing the model architecture is important for:

• Model explainability
• Results interpretation
• Model debugging
• Model comparison

Read more about model visualization

Visualizing Machine Learning Models: Guide

Once you get some decent understanding of one model, you are good, right? Wrong 🙂

Typically, you need to do some or a lot of experimenting with model improvement ideas, and visualizing differences between various ML experiments becomes crucial. 

There are many methods that you can use to get that understanding:

• Look at evaluation metrics (also, you should know how to choose an evaluation metric for your problem)
• Look at performance charts like ROC, Lift Curve, Confusion Matrix, and others
• Look at learning curves to estimate overfitting
• Look at model predictions on best/worst cases
• Look how resource-intensive is model training and inference (they translate to serious costs and will be crucial to the business side of things) 
• Use model interpretation tools and techniques to vet predictions (further reading)

You can do all of those (or most of those) yourself, but today there are tools that you can use. If you’re looking for the best tools that will help you visualize machine learning experiments and models, you’ve come to the right place.

Tools for Machine Learning experiments visualization

1. Neptune

Neptune is a metadata store for MLOps, built for teams that run a lot of experiments.‌ It gives you a single place to log, store, display, organize, compare, and query all your model-building metadata.

Neptune is used for:‌

• Experiment tracking: Log, display, organize, and compare ML experiments in a single place.
• Model registry: Version, store, manage, and query trained models and model building metadata.
• Monitoring ML runs live: Record and monitor model training, evaluation, or production runs live.

How can Neptune help you visualize experiments and models?

• Log any metadata type, including parameters and metrics, but also rich objects like images, video, audio, and interactive visualizations.
• Visualize logged metadata and analyze results in a preferred way – in the runs table, as charts, in dashboards, or in a folder structure.
• Compare hyperparameters and metrics across many runs with an intelligent compare table that highlights what was different.
• See how different parameters and configs affect the results, and debug and optimize models.
• Automatically monitor hardware utilization (GPU, CPU, memory).
• Register models and metadata associated with them.

Learn more

Play with a live example in the Neptune app (no registration needed).

2. Weights & Biases

Weights & Biases is a machine learning platform for developers to build better models faster. It lets you quickly track experiments, version and iterate on datasets, evaluate model performance, reproduce models, visualize results and spot regressions, and share findings with colleagues.

How can Weights & Biases help you visualize experiments and models?

• Monitor training runs information like loss, accuracy (learning curves)
• View histograms of weights and biases (no pun intended), or gradients
• Log rich objects like, charts, video, audio or interactive charts during training
• Use various comparison tools like tables showing auto-diffs, parallel coordinates plot and others
• Interactive prediction bounding box visualization for object detection models
• Interactive prediction masks visualization for semantic segmentation models
• Visualize live metrics like GPU and CPU utilization
• Build dataset dependency graphs
• Visualize parameter importance

Read more

See an in-depth comparison between Neptune and Weights & Biases.

3. Comet

Comet is a meta machine learning platform for tracking, comparing, explaining, and optimizing experiments and models. It also lets you monitor, detect and alert, and debug production models in real time. You can also build your own or use community provided custom visualizations and applications based on experiments and model data.

How can Comet help you visualize experiments and models?

• Visualize samples with dedicated modules for vision, audio, text and tabular data to detect overfitting and easily identify issues with your dataset
• You can customize and combine your visualizations
• You can monitor your learning curves
• Comet’s flexible experiments and visualization suite allow you to record, compare, and visualize many artifact types
• Build your own, or use community built, “panels” to visualize your models and experiments
• Monitor production models in real-time

Read more

See an in-depth comparison between Neptune and Comet.

4. TensorBoard

TensorBoard provides the visualization and tooling needed for machine learning experimentation. It’s open-source and offers a suite of tools for visualization and debugging of machine learning models. TensorBoard is the most popular solution on the market; thus, it’s widely integrated with many other tools and applications.

What’s more, it has an extensive network of engineers using this software and sharing their experience and ideas. This makes a powerful community ready to solve any problem. The software, itself, however, is best suited for an individual user.

How can TensorBoard help you visualize experiments and models?

• Tracking and visualizing metrics such as loss and accuracy
• Visualizing the model graph (ops and layers)
• Viewing histograms of weights, biases, or other tensors as they change over time
• Projecting embeddings to a lower-dimensional space
• Displaying images, text, and audio data
• Profiling TensorFlow programs

5. Sacred + Omniboard

The Sacred + Omniboard stack actually comprises two open-source components – Sacred, an experiment management tool, and Omniboard, a web dashboard for Sacred. Sacred logs experiment data to a MongoDB backend which Omniboard connects to visualize experiments.

Sacred gives you a powerful command-line interface, observers to log experiment information, and automatic seeding to control experiment randomness so that results are reproducible. Omniboard adds to this with its suite of experiment management, drill down, and comparison functionalities.

How can Sacred + Omniboard help you visualize experiments and models?

• The Sacred CLI can be used to update parameters and run different variants of the same experiment
• Save experiment configurations in a MongoDB database
• Compare different experiments on the Omniboard dashboard
• Roll-up or drill-down experiments to either view aggregates or individual details

6. MLflow

MLflow is an open-source platform for managing the end-to-end machine learning lifecycle. It is currently by far the most widely used ML experiment tracker and as such, is widely supported by languages, frameworks, and platforms. MLflow supports logging through Python, R, Java, and REST APIs.

Being an open-source project, there are no direct costs involved with MLflow for personal use. However, infrastructure costs will add up if you want to use MLflow for collaboration between large teams, as you will need to have your own remote tracking servers. Databricks’ managed MLflow takes care of hosting for you.

How can MLflow help you visualize experiments and models?

• Log parameters, code versions, metrics, and artifacts
• Log from a script or notebook
• Log to a locally or remotely hosted server
• Visualize metrics
• Organize runs into experiments
• Automatic logging with some of the most widely used ML frameworks

Comparison table: tools for ML experiments visualization

Experiment tracking is the core use-case of all the abovementioned products, and as such, most offer similar functionality. However, a few features are not supported in all the tools. The below table highlights a few differences in the functionality offered by the abovementioned tools.

Machine Learning model visualization tools

1. dtreeviz

dtreeviz is a python library for decision tree visualization and model interpretation. It currently supports scikit-learn, XGBoost, Spark MLlib, and LightGBM trees. Versions 1.3 onwards also support feature space illustrations for any model which have a predict_proba().

How can dtreevix help you visualize models?

• Visualize tree-based models
• Illustrate feature space for supported classifiers
• Highlight prediction path for a single observation in the tree, and get plain English explanations for the same
• Visualize leaf metrics like purity, samples, and target distribution
• Visualize classification boundaries in the feature space.

2. Netron

Netron can be used to produce interactive visualizations of neural network models. It is available as a Python package, a standalone application, and a browser application.

How can Netron help you visualize models?

• Produce interactive visualizations of neural networks 
• Expand nodes to view layer details
• Supports almost all widely used frameworks
• No installation required to use the browser application
  

3. NN-SVG

NN-SVG is a hosted application for creating Neural Network (NN) architecture drawings parametrically rather than manually. It also provides the ability to export those drawings to Scalable Vector Graphics (SVG) files, suitable for inclusion in academic papers or web pages.

How can NN-SVG help you visualize models?

• Fully hosted application, no installation required
• Supports FCNN, LeNet, and AlexNet styles
• Highly customizable
• Export generated schematics as an SVG image

4. TensorBoard

TensorBoard is mostly an experiment visualization tool, but it also shows model graphs for TensorFlow models.

How can TensorBoard help you visualize models?

• Visualize TensorFlow model graphs
• Track and log TensorFlow experiments
• Provides a fully-managed version – TensorBoard.dev – which lets you host, track, and share your ML experiments 

Summary

In this blog, we have covered and compared some of the most popular model and visualization tools available currently. As with most of the tools in the ML space, there is no one-size-fits-all, and which tool to use will depend a lot on your specific use case. 

What kind of models do you want to visualize? And what is the purpose of your visualizations?

Most of the model visualization tools are built targeting specific architectures and use cases.

• For tree-based models, dtreeviz is a good choice.
• TensorBoard integrates very well with Tensorflow/Keras models.
• If you just want to create a non-interactive visualization of your neural network model for illustration purposes, NN-SVG might save you some time.
• Netron produces detailed interactive visualizations of neural network models, which can help a lot with debugging and optimization. 

Are you a small team working on a limited budget?

If yes, some of the open-source and free tools might be good to start with. Keep in mind, though, that as your team grows, your self-hosting costs will grow as well. You may also consider SaaS solutions with pricing that doesn’t depend on the number of users (e.g., Neptune). This way, you’ll be better prepared for the time when your needs scale.

Are you a larger team looking for collaborative development?

Most SaaS solutions have rich collaboration features like URL-based sharing, collaborative dashboards, and reports, user management, etc. 

In any case, hope this article will help you make that choice, and you’ll find what you were looking for.

Happy visualizing!