Using MLflow with ATOM to track all your machine learning experiments without additional code

Author:Murphy | View: 25286 | Time: 2025-03-23 19:22:20

Introduction

The MLflow Tracking component is an API and UI for logging parameters, code versions, metrics, and output files when running your Machine Learning experiments and for later visualizing the results.

In this story, we'll explain how to use the ATOM library to easily track your models, parameters, pipelines, data and plots. ATOM is an open-source Python package designed to help data scientists with the exploration of machine learning pipelines.

Note: This story focuses on using the experiment tracking feature of ATOM. Teaching the basics of the library lies outside the scope of this story. Read this other story if you want a gentle introduction to the library.

Feature overview

Start tracking your experiments assigning a name to the experiment parameter in atom's constructor. Every model is tracked using a separate run. When no backend is configured, the data is stored locally at ./mlruns. To configure the backend, use mlflow.set_tracking_uri in your notebook or IDE before initializing atom. This does not affect the currently active run (if one exists), but takes effect for successive runs. Run Mlflow ui on your terminal to open MLflow's Tracking UI and view it at http://localhost:5000.

Note: When using ATOM on Databricks, the experiment's name should include the complete path to the storage, e.g. /Users/[email protected]/experiment_name.

The following elements are tracked:

Tags The runs are automatically tagged with the model's full name, the branch from which the model was trained, and the time it took to fit the model. Add additional custom tags through the ht_params parameter, e.g. atom.run(["LR", "RF"], ht_params={"tags": {"tag1": 1}}).

Parameters All parameters used by the estimator at initialization are tracked. Additional parameters passed to the fit method are not tracked.

Model The model's estimator is stored as artifact. This option can be switched off using atom's [log_model](https://tvdboom.github.io/ATOM/v5.1/API/ATOM/atomclassifier/#atomclassifier-log_model) attribute, e.g. atom.log_model = False.

Hyperparameter tuning If hyperparameter tuning is performed, every trial is tracked as a nested run in the model's main run. This option can be switched off using atom's [log_ht](https://tvdboom.github.io/ATOM/v5.1/API/ATOM/atomclassifier/#atomclassifier-log_ht) attribute, e.g. atom.log_ht = False.

Metrics All metric results are tracked, not only during training, but also when the evaluate method is called at a later point. Metrics calculated during in-training validation are also stored.

Dataset The train and test sets used to fit and evaluate the model can be stored as .csv files to the run's artifacts. This option can be switched on using atom's [log_data](https://tvdboom.github.io/ATOM/v5.1/API/ATOM/atomclassifier/#atomclassifier-log_data) attribute, e.g. atom.log_data = True.

Pipeline The model's pipeline (returned from the export_pipeline method) can be stored as an artifact. This option can be switched on using atom's [log_pipeline](https://tvdboom.github.io/ATOM/v5.1/API/ATOM/atomclassifier/#atomclassifier-log_pipeline) attribute, e.g. atom.log_pipeline = True.

Plots By default, plots are stored as .html artifacts in all runs corresponding to the models that are showed in the plot. If the filename parameter is specified, they are stored under that name, else the method's name is used. This option can be switched off using atom's [log_plots](https://tvdboom.github.io/ATOM/v5.1/API/ATOM/atomclassifier/#atomclassifier-log_plots) attribute, e.g. atom.log_plots = False.

Example

The easiest way to fully understand the capabilities is through an example. We initialize atom the usual way, and specify the experiment parameter. The name provided here is the name of the mlflow experiment used. If there is no experiment with that name, a new one is created.

from atom import ATOMClassifier
from sklearn.datasets import load_breast_cancer

X, y = load_breast_cancer(return_X_y=True, as_frame=True)

atom = ATOMClassifier(X, y, experiment="breast_cancer")

Let's specify that we want to log the data and pipeline as well. These two options are turned off by default.

atom.log_data = True
atom.log_pipeline = True

Train the models the usual way. All selected metrics are logged.

atom.run(models=["LR", "RF"], metric=["f1", "accuracy", "precision"])

Now open the UI running mlflow ui in the terminal. Every model has its separate run.

Per run, atom stores:

The model's parameters:

The metric scores on train and test set:

Predefined tags and custom tags (if specified):

Artifacts, which include the estimator and pipeline as pickle files, and the train and test sets as csv files.

Additional metrics (calculated with the evaluate method) and plots where the model is displayed are also added to the run.

atom.evaluate()

atom.plot_gains()

Hyperparameter tuning

When running hyperparameter tuning, every trial of the study is added as nested run to the main run.

atom.run(models="LGB", n_trials=10)

The nested runs are named - . Their metric score is not on train nor test, but on the validation set instead, which is the subset of the training set used to validate this specific trial (the rest of the set was used to fit the estimator). The test set is not used during hyperparameter tuning to avoid data leakage.

Note: The data and pipeline are never stored within nested runs.

In-training validation

Some of the models allow in-training validation. This means that the estimator is evaluated (using only the first metric) on the train and test set after every round of the training (a round can be an iteration for linear models or an added tree for boosted tree models). The validation scores are stored in the evals metric, which is also tracked in mlflow.

Note: The evals metric is not calculated in nested runs.

DAGsHub integration

ATOM has a build-in integration with DAGsHub, a web platform based on open source tools, optimized for data science and oriented towards the open source community. To store your mlflow experiments in a DAGsHub repo, type dagshub: in the experiment parameter (instead of just the experiment's name). If the repo does not already exist, a new public repo is created.

A minimalistic example could look like this:

from atom import ATOMClassifier
from sklearn.datasets import load_breast_cancer

X, y = load_breast_cancer(return_X_y=True, as_frame=True)

atom = ATOMClassifier(X, y, experiment="dagshub:breast_cancer")

atom.run(models=["LR", "RF"])

Note: If you are logged into your DAGsHub account when running atom's constructor, a page on your web browser is automatically opened to give access permissions. If not, read here how to set up your DAGsHub credentials.

Conclusion

We have shown how easy it is to track your machine learning experiments using the ATOM library. With minimal code changes, it's now possible to store the estimator, parameters, pipelines, data and plots for every model trained.

For further information about ATOM, have a look at the package's documentation. For bugs or feature requests, don't hesitate to open an issue on GitHub or send me an email.