Track Computer Vision Experiments with MLflow

Discover how to set up an efficient MLflow environment to track your experiments, compare and choose the best model for deployment

Training and fine-tuning various models is a basic task for every computer vision researcher. Even for easy ones, we do a hyper-parameter search to find the optimal way of training the model over our custom dataset. Data augmentation techniques (which include many different options already), the choice of optimizer, learning rate, and the model itself. Is it the best architecture for my case? Should I add more layers, change the architecture, and many more questions will wait to be asked and searched?

While searching for an answer to all these questions, I used to save the model training process log files and output checkpoints in different folders in my local, change the output directory name every time I ran a training, and compare the final metrics manually one-by-one. Tackling the experiment-tracking process in such a manual way has many disadvantages: it’s old school, time and energy-consuming, and prone to errors.

In this blog post, I will show you how to use MLflow, one of the best tools to track your experiment, allowing you to log whatever information you need, visualize and compare the different training experiments you have accomplished, and decide which training is the optimal choice in a user- (and eyes-) friendly environment!