ML Experiment Tracking: Everything You Need to Know

• use different models and model hyperparameters,
• use different training or evaluation data, 
• run different code (including this small change that you wanted to test quickly),
• run the same code in a different environment (not knowing which PyTorch or Tensorflow version was installed)

In this article, you will learn:

• What is ML experiment tracking
• 4 ways in which it can improve your work
• What are the best practices of ML experiment tracking
• How to add experiment tracking into your workflow

What is ML experiment tracking?

Experiment tracking is the process of saving all experiment related information that you care about for every experiment you run.   

• Scripts used for running the experiment
• Environment configuration files 
• Versions of the data used for training and evaluation
• Parameter configurations
• Evaluation metrics 
• Model weights
• Performance visualizations (confusion matrix, ROC curve)  
• Example predictions on the validation set (common in computer vision)

• Experiment database: A place where experiment metadata is stored and can be logged and queried
• Experiment dashboard: A visual interface to your experiment database. A place where you can see your experiment metadata.
• Client library: Which gives you methods for logging and querying data from the experiment database. 

Experiment tracking vs ML model management vs MLOps

• streamlines moving models from experimentation to production
• helps with model versioning
• organizes model artifacts in an ML model registry
• helps with testing various model versions in the production environment
• enables rolling back to an old model version if the new one seems to be going crazy

Experiment tracking is useful even if your models don’t make it to production (yet).

Why does experiment tracking matter? 

All of your ML experiments are organized in a single place

• Private laptop
• PC at work
• A dedicated instance in the Cloud 
• University cluster
• Kaggle kernel or Google Colab
• And many more. 

“[experiment tracking system] allows us to keep all of our experiments organized in a single space. Being able to see my team’s work results any time I need makes it effortless to track progress and enables easier coordination.” – Michael Ulin VP, Machine Learning @Zesty.ai

• Search and filter experiments to find the information you need quickly
• Compare their metrics and parameters with no additional work
• Drill down and see what exactly it was that you tried (code, data versions, architectures) 
• Reproduce or re-run experiments when you need to
• Access experiment metadata even if you don’t have access to the server where you ran them 

Compare experiments, analyze results, debug model training with little extra work

• the way you log things can change, 
• you may forget to log something important,
• you may simply lose some information accidentally. 

“Tracking and comparing different approaches has noticeably boosted our productivity, allowing us to focus more on the experiments [and] develop new, good practices within our team…” – Tomasz Grygiel, Data Scientist @idenTT

• Compare parameters and metrics
• Overlay learning curves
• Group and compare experiments based on data versions or parameter values
• Compare Confusion Matrices, ROC curves, or other performance charts
• Compare best/worst predictions on test or validation sets
• View code diffs (and/or notebook diffs)
• Look at hardware consumption during training runs for various models
• Look at prediction explanations like Feature Importance, SHAP or Lime
• Compare rich-format artifacts like video or audio 
• …Compare anything else you logged 

Improve collaboration: see what everyone is doing, share experiment results easily, access experiment data programmatically   

“[An experiment tracking system] makes it easy to share results with my teammates. I’m sending them a link and telling what to look at, or I’m building a view on the experiments dashboard. I don’t need to generate it by myself, and everyone in my team has access to it.” – Maciej Bartczak, Resarch Lead @Banacha Street

See your ML runs live: manage experiments from anywhere and anytime

“Without information I have in the monitoring section I wouldn’t know that my experiments are running 10 times slower than they could.” – Michał Kardas, Machine Learning Researcher @TensorCell

Experiment tracking best practices

What you should keep track of in any ML experiment

• Code: preprocessing, training and evaluation scripts, notebooks used for designing features, other utilities. All the code that is needed to run (and re-run) the experiment.  
• Environment: The easiest way to keep track of the environment is to save the environment configuration files like `Dockerfile` (Docker), `requirements.txt` (pip) or `conda.yml` (conda). You can also save the Docker image on Docker Hub, but I find saving configuration files easier.   
• Data: saving data versions (as a hash or locations to data files) makes it easy to see what your model was trained on. You can also use modern data versioning tools like DVC (and save the .dvc files to your experiment tracking tool). 
• Parameters: saving your run configuration is absolutely crucial. Be especially careful when you pass parameters via the command line (argparse, click, hydra) as this is a place where you can easily forget to track (I have some horror stories to share). hyperparameters. 
• Metrics: logging evaluation metrics on train, validation, and test sets for every run is pretty obvious. But different frameworks do it differently.

• Model weights
• Evaluation charts (ROC Curve, Confusion matrix)
• Prediction distributions

• Model checkpoints (both during and after training)
• Gradient norms (to control for vanishing or exploding gradient problems) 
• Best/worst predictions on validation/test set after training
• Hardware resources: especially useful in debugging data loaders and multi GPU setups

• Model Predictions after every epoch (labels, overlayed masks or bounding boxes)

• Prediction explanations (eli5 text explainer is good) on evaluation/test data

• Input data snapshot ( `.head()` on training data if you are using pandas)
• Feature importances (permutation importance)
• Prediction explanations like SHAP or partial dependence plots (they are all available in DALEX). 

• Episode return and episode length
• Total environment steps, wall time, steps per second
• Value and police function losses
• Aggregate statistics over multiple environments and/or runs

• Run score: metric you are optimizing with HPO after every iteration
• Run parameters: parameters tried at each iteration
• Best parameters: best parameters so far and best parameters after the HPO sweep is finished
• Parameter comparison charts: there are various visualizations that you may want to log during or after training, like parallel coordinates plot or slice plot (they are all available in Optuna by the way).

How to set up experiment tracking

• Spreadsheets + naming conventions
• Versioning configuration files with Github
• Using modern experiment tracking tools 

You can use Spreadsheets and naming conventions (but please don’t)

…things can fall apart really quickly

• You have to remember to track them. If something doesn’t happen automatically, things get messy, especially with more people involved.
• You have to be sure that you or your team will not overwrite things in the spreadsheet by accident. Spreadsheets are not easy to version, so if this happens, you are in trouble. 
• You have to remember to use the naming conventions. If someone on your team messes this up, you may not know where the experiment artifacts (model weights, performance charts) for the experiments you ran are. 
• You have to back up your artifact directories (remember that things break). 
• When your spreadsheet grows, it becomes less and less usable. Searching for things and comparing hundreds of experiments in a spreadsheet (especially if you have multiple people that want to use it at the same time) is not a great experience.  

You can version metadata files in GitHub

… Github wasn’t built for … machine learning

• .git and Github wasn’t built for comparing machine learning objects. 
• Comparing more than two experiments is not going to work. Compare in .git systems was designed for comparing two branches, master and develop, for example. If you want to compare multiple experiments, take a look at metrics and overlay learning curves you are out of luck. 
• Organizing many experiments is difficult (if not impossible). You can have branches with ideas or a branch per experiment but the more experiments you run the less usable it becomes 
• You will not be able to monitor your experiments live, the information will be saved after your experiment is finished. 

You can use one of the modern experiment tracking tools

“Within the first few tens of runs, I realized how complete the tracking was – not just one or two numbers, but also the exact state of the code, the best-quality model snapshot stored to the cloud, the ability to quickly add notes on a particular experiment. My old methods were such a mess by comparison.” – Edward Dixon, Data Scientist @intel

import neptune.new as neptune
run = neptune.init(...) # create a Run credentials

from neptune.new.types import File
run['accuracy'] = 0.92
for prediction_image in worst_predictions:
    run['worst predictions'].log(File.as_image(prediction_image))

python train.py

• be easier to use for a machine learning person than general tools
• have better integrations with the ML ecosystem
• have more experiment-focused features than the general solutions  

Next steps

• what to keep track of for your project, 
• how to use that information to improve future experiments,
• how to improve your teams’ unique workflow with it,
• when to even use experiment tracking.

• Create a free Neptune account
• Try Neptune on Colab (zero setup, no registration)

• Article: A Complete Guide to Monitoring ML Experiments Live in Neptune
• Article: How to Set Up Continuous Integration for Machine Learning with Github Actions and Neptune