Experiment Tracking with MLflow and Dagshub: A Complete Guide

Introduction

In the world of machine learning, reproducibility and organization are not just nice-to-have features — they’re essential components of successful projects. As data scientists, we run numerous experiments with different parameters, datasets, and models. Without proper tracking, this can quickly turn into chaos.

Enter MLflow and DAGsHub — two powerful tools that, when combined, create a robust experiment tracking system that can transform your ML workflow.

What You’ll Learn

• The importance of experiment tracking in ML
• How to set up MLflow with DAGsHub
• Tracking experiments with code examples
• Analyzing and comparing results
• Best practices for ML experiment management

Why Experiment Tracking Matters

Before we dive into the technical implementation, let’s understand why experiment tracking is crucial:

1. Reproducibility: Recreate any model exactly as it was trained
2. Comparison: Systematically compare different approaches
3. Collaboration: Share results with team members effectively
4. Debugging: Identify what went wrong in failed experiments
5. Compliance: Maintain audit trails for regulated industries

Setting Up Your Environment

Prerequisites

# Create a new conda environment (optional but recommended)
conda create -n mlflow-tracking python=3.12
conda activate mlflow-tracking

# Install required packages
pip install mlflow dagshub scikit-learn pandas numpy matplotlib

DAGsHub Account Setup

1. Create a DAGsHub Account:

• Visit DAGsHub
• Sign up for a free account
• Create a new repository

2. Get Your DAGsHub Credentials:

• Go to your profile settings
• Generate an access token (Settings → Access Tokens)

Configuration

import os
import mlflow
import dagshub

# Configure DAGsHub
DAGSHUB_USERNAME = "your_username"
DAGSHUB_REPO_NAME = "your_repo_name"
DAGSHUB_TOKEN = "your_access_token"

# Initialize DAGsHub connection
dagshub.init(repo_owner=DAGSHUB_USERNAME, 
             repo_name=DAGSHUB_REPO_NAME,
             mlflow=True)

# Set tracking URI
mlflow.set_tracking_uri(f"https://dagshub.com/{DAGSHUB_USERNAME}/{DAGSHUB_REPO_NAME}.mlflow")

Basic Experiment Tracking with MLflow

Let’s start with a simple example using a scikit-learn classifier.

Example 1: Basic Classification Experiment

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from sklearn.datasets import load_iris

# Load dataset
data = load_iris()
X = pd.DataFrame(data.data, columns=data.feature_names)
y = data.target

# Split the data
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

def run_experiment(params):
    """Run a single experiment with given parameters"""

    # Start MLflow run
    with mlflow.start_run():
        # Log parameters
        mlflow.log_params(params)

        # Create and train model
        model = RandomForestClassifier(
            n_estimators=params['n_estimators'],
            max_depth=params['max_depth'],
            random_state=42
        )

        model.fit(X_train, y_train)

        # Make predictions
        y_pred = model.predict(X_test)

        # Calculate metrics
        accuracy = accuracy_score(y_test, y_pred)
        precision = precision_score(y_test, y_pred, average='weighted')
        recall = recall_score(y_test, y_pred, average='weighted')
        f1 = f1_score(y_test, y_pred, average='weighted')

        # Log metrics
        mlflow.log_metrics({
            'accuracy': accuracy,
            'precision': precision,
            'recall': recall,
            'f1_score': f1
        })

        # Log model
        mlflow.sklearn.log_model(model, "random_forest_model")

        # Log artifacts (example: feature importance plot)
        import matplotlib.pyplot as plt

        plt.figure(figsize=(10, 6))
        feature_importances = pd.Series(model.feature_importances_, index=X.columns)
        feature_importances.nlargest(10).plot(kind='barh')
        plt.title('Feature Importances')
        plt.tight_layout()

        # Save plot
        plt.savefig('feature_importance.png')
        mlflow.log_artifact('feature_importance.png')

        print(f"Experiment completed with accuracy: {accuracy:.4f}")

# Run multiple experiments with different parameters
experiment_params = [
    {'n_estimators': 50, 'max_depth': 3},
    {'n_estimators': 100, 'max_depth': 5},
    {'n_estimators': 200, 'max_depth': 7},
    {'n_estimators': 300, 'max_depth': None},
]

for params in experiment_params:
    run_experiment(params)

Advanced Experiment Tracking

Custom Metrics and Artifacts

import json
from datetime import datetime

def run_advanced_experiment():
    """Example of advanced experiment tracking"""

    with mlflow.start_run(run_name=f"advanced_experiment_{datetime.now().strftime('%Y%m%d_%H%M%S')}"):

        # Parameters
        params = {
            'n_estimators': 150,
            'max_depth': 10,
            'min_samples_split': 2,
            'min_samples_leaf': 1,
            'bootstrap': True
        }

        mlflow.log_params(params)

        # Train model
        model = RandomForestClassifier(**params, random_state=42)
        model.fit(X_train, y_train)

        # Predictions
        y_pred = model.predict(X_test)
        y_pred_proba = model.predict_proba(X_test)

        # Comprehensive metrics
        from sklearn.metrics import classification_report, confusion_matrix

        accuracy = accuracy_score(y_test, y_pred)
        mlflow.log_metric('accuracy', accuracy)

        # Log classification report as JSON
        clf_report = classification_report(y_test, y_pred, output_dict=True)
        with open('classification_report.json', 'w') as f:
            json.dump(clf_report, f, indent=2)
        mlflow.log_artifact('classification_report.json')

        # Log confusion matrix as plot
        import seaborn as sns
        plt.figure(figsize=(8, 6))
        cm = confusion_matrix(y_test, y_pred)
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
        plt.title('Confusion Matrix')
        plt.ylabel('True Label')
        plt.xlabel('Predicted Label')
        plt.savefig('confusion_matrix.png')
        mlflow.log_artifact('confusion_matrix.png')

        # Log training dataset info
        dataset_info = {
            'training_samples': len(X_train),
            'test_samples': len(X_test),
            'features': list(X.columns),
            'target_classes': list(data.target_names)
        }
        mlflow.log_dict(dataset_info, 'dataset_info.json')

        # Set tags
        mlflow.set_tag('model_type', 'RandomForest')
        mlflow.set_tag('dataset', 'Iris')
        mlflow.set_tag('author', 'Your Name')

        print(f"Advanced experiment completed. Accuracy: {accuracy:.4f}")

run_advanced_experiment()

Hyperparameter Tuning with MLflow Tracking

from sklearn.model_selection import GridSearchCV

def hyperparameter_tuning():
    """Example of tracking hyperparameter tuning"""

    with mlflow.start_run(run_name="hyperparameter_tuning"):

        # Define parameter grid
        param_grid = {
            'n_estimators': [50, 100, 200],
            'max_depth': [3, 5, 7, None],
            'min_samples_split': [2, 5, 10],
            'min_samples_leaf': [1, 2, 4]
        }

        # Setup grid search
        grid_search = GridSearchCV(
            RandomForestClassifier(random_state=42),
            param_grid,
            cv=5,
            scoring='accuracy',
            n_jobs=-1,
            verbose=1
        )

        # Perform grid search
        grid_search.fit(X_train, y_train)

        # Log best parameters and score
        mlflow.log_params(grid_search.best_params_)
        mlflow.log_metric('best_cv_score', grid_search.best_score_)
        mlflow.log_metric('test_score', grid_search.score(X_test, y_test))

        # Log all CV results
        cv_results = pd.DataFrame(grid_search.cv_results_)
        cv_results.to_csv('cv_results.csv', index=False)
        mlflow.log_artifact('cv_results.csv')

        # Log best model
        mlflow.sklearn.log_model(grid_search.best_estimator_, "best_grid_search_model")

        print(f"Best parameters: {grid_search.best_params_}")
        print(f"Best CV score: {grid_search.best_score_:.4f}")

        return grid_search.best_estimator_

best_model = hyperparameter_tuning()

Viewing and Analyzing Results on DAGsHub

After running your experiments, you can view them on DAGsHub:

1. Navigate to your DAGsHub repository
2. Click on the “Experiments” tab
3. View all your runs with their parameters and metrics

Comparing Experiments

DAGsHub provides an intuitive interface to:

• Sort experiments by different metrics
• Filter runs based on parameters or tags
• Compare multiple runs side by side
• Visualize metrics and parameters

Best Practices for Experiment Tracking

1. Consistent Naming Conventions

# Good practice: Descriptive run names
mlflow.start_run(run_name="rf_200_trees_7_depth_adam_optimizer")

# Bad practice: Non-descriptive names
mlflow.start_run(run_name="experiment_1")

2. Comprehensive Logging

# Always log:
# - All hyperparameters
# - Key metrics (multiple, not just accuracy)
# - Dataset version or characteristics
# - Environment details (optional but helpful)
# - Model artifacts
# - Visualizations and plots

3. Organize with Tags

mlflow.set_tag('project', 'customer_churn_prediction')
mlflow.set_tag('phase', 'exploration')
mlflow.set_tag('data_version', 'v2.1')
mlflow.set_tag('priority', 'high')

4. Version Control Integration

# Commit your code alongside experiments
git add .
git commit -m "Add hyperparameter tuning experiment"
git push

Complete Project Structure

mlflow-dagshub-project/
│
├── data/
│   └── iris.csv
├── notebooks/
│   └── exploration.ipynb
├── src/
│   ├── __init__.py
│   ├── data_processing.py
│   ├── model_training.py
│   └── utils.py
├── experiments/
│   └── run_experiments.py
├── requirements.txt
└── README.md

Troubleshooting Common Issues

Connection Problems

# Check your connection
try:
    dagshub.init(repo_owner=DAGSHUB_USERNAME, 
                 repo_name=DAGSHUB_REPO_NAME,
                 mlflow=True)
    print("Connection successful!")
except Exception as e:
    print(f"Connection failed: {e}")

Authentication Issues

# Make sure your token is correct
import getpass

if DAGSHUB_TOKEN == "your_access_token":
    print("Please set your actual DAGsHub token!")
    # Optionally prompt for token
    DAGSHUB_TOKEN = getpass.getpass("Enter your DAGsHub token: ")

Conclusion

Experiment tracking with MLflow and DAGsHub provides a powerful, scalable solution for managing your machine learning workflows. By implementing the practices outlined in this guide, you’ll be able to:

• ✅ Track all experiments systematically
• ✅ Reproduce results reliably
• ✅ Collaborate effectively with team members
• ✅ Make data-driven decisions about model improvements
• ✅ Maintain organized records of your ML projects

The combination of MLflow’s robust tracking capabilities with DAGsHub’s user-friendly interface creates an ecosystem where you can focus on what matters most — building better models.

Next Steps

1. Set up your own DAGsHub repository
2. Try the code examples with your own datasets
3. Explore advanced features like model registry
4. Integrate experiment tracking into your existing projects

Additional Resources

[embed]MLflow MLflow Documentation - Machine Learning and GenAI lifecycle managementmlflow.org

[embed]DagsHub Documentation | DagsHub Docs Learn how to use DagsHub to build multimodal AI models and datasets in a single platform. Curate and annotate datasets…dagshub.com

[embed]| DagsHub Docs DagsHub Documentation - Learn how to use DagsHub and improve your data science workflowdagshub.com

This post demonstrates the power of proper experiment tracking. Start implementing these practices in your projects today, and you’ll never lose track of your experiments again!

Happy experimenting! 🚀

About the Author: Sayantan Das is a data scientist passionate about AI and ML engineering best practices and reproducible research. Connect with me on

https://www.linkedin.com/in/sayantandas1989/

If you found this article helpful, please clap 👏 and share it with your colleagues!