Tracking ML Experiments With Data Version Control

Machine Learning Operations, or MLOps, is a practice for collaboration and communication between data scientists and operations professionals to help manage the production ML lifecycle. One key aspect of MLOps is experiment tracking – recording the different models, parameters, features, and results of each trial in order to identify the best performing model. However, the modern machine learning lifecycle generates a huge volume of metadata that can be challenging to manage – data, model versions, hyperparameters, metrics, and so on.

This is where data version control comes in. Just as Git enables source control for code, data version control allows data scientists to track their ML experiments in a reproducible way. By versioning data and trained models along with the code, it provides a single source of truth and makes it easy to collaborate with others. Let‘s dive deeper into data version control and see how tools like DVC can help streamline your MLOps workflow.

What is Data Version Control?

Data Version Control, or DVC, is an open-source tool for tracking and versioning machine learning models and datasets. It works by storing information about your data and model files in a special .dvc file which is checked into Git alongside your code. These .dvc files act as a pointer to the actual data or model file, which can be stored on a local file system, remote storage like S3, or even an HTTP server.

With DVC, you can version your data and models just like you version your code with Git. This means every time you run an experiment, the input data, configuration parameters, and output models and metrics are captured in a .dvc file. You can then use Git to create a version history of your project, tracking which data and model was used for each experiment.

The beauty of this approach is that it makes your experiments reproducible. By checking out a specific Git commit, you can pull the exact data and model used in that experiment and rerun it. DVC also makes it easy to switch between experiments – just check out the desired Git commit and run dvc repro to execute the pipeline with those specific versions of data and models.

There are several open source data version control tools available, including:

• DVC: A Git-compatible data version control tool
• ML-Metadata: A library for recording and retrieving metadata associated with ML workflows
• Pachyderm: A data versioning and pipelines platform based on Kubernetes and Git

In this post we‘ll focus on DVC, but the concepts apply to other tools as well. The key idea is to decouple code and data, but keep them tightly linked through .dvc files checked into Git.

Example Project

To demonstrate how data version control works in practice, let‘s walk through an example machine learning project. We‘ll use the classic Iris dataset to build a model that predicts the species of a flower based on measurements of its petals and sepals.

The project will have the following structure:

iris/
    data/
        raw/
            iris.csv
    processed/
        train.csv
        test.csv
    src/
        prepare.py
        train.py
        evaluate.py
    dvc.yaml        
    params.yaml

The raw/ directory contains the original Iris dataset. The processed/ directory will contain the data after it has been split into train and test sets. The src/ directory has Python scripts for each stage in the ML pipeline – preparing the data, training the model, and evaluating its performance. The dvc.yaml file specifies the pipeline stages and their dependencies, while the params.yaml file contains configuration parameters like the train/test split ratio and model hyperparameters.

To get started, we initialize Git and DVC in the project directory:

git init
dvc init

Next we configure the remote storage where DVC will store the actual data and model files. For this example we‘ll use a local remote:

dvc remote add -d storage /tmp/dvc-storage
git commit .dvc/config -m "Configure local remote"

Now let‘s add the raw Iris dataset to DVC:

dvc add data/raw
git add data/raw.csv.dvc
git commit -m "Add raw data"

This creates a data/raw.csv.dvc file that points to the actual data file in the /tmp/dvc-storage directory. We add this .dvc file to Git so it can be versioned alongside the code.

With the data in place, we can define the pipeline stages in the dvc.yaml file:

stages:
  prepare:
    cmd: python src/prepare.py
    deps:
      - data/raw/iris.csv
      - src/prepare.py
    params:
      - prepare.train_size
      - prepare.random_state
    outs:
      - data/processed/train.csv
      - data/processed/test.csv
  train:
    cmd: python src/train.py
    deps:
      - data/processed/train.csv
      - src/train.py 
    params:
      - train.n_estimators
      - train.min_samples_split
    outs:
      - model.pkl
  evaluate:
    cmd: python src/evaluate.py
    deps:
      - data/processed/test.csv
      - model.pkl
      - src/evaluate.py
    metrics:
      - metrics.json

This defines three stages – prepare, train, and evaluate – each of which executes a Python script. The prepare stage takes the raw Iris data and splits it into train/test sets, while the train stage trains a Random Forest model on the training set. Finally, the evaluate stage calculates performance metrics on the test set.

Each stage lists its dependencies (deps) – the data and code files it needs as inputs – and the outputs (outs) it produces. The prepare and train stages also specify parameters that are read from the params.yaml file:

prepare:
  train_size: 0.8
  random_state: 42

train:
  n_estimators: 100
  min_samples_split: 2

To execute the pipeline, we simply run:

dvc repro 

DVC will execute each stage in the correct order based on the dependencies specified in dvc.yaml. The output data and model files will be saved in the /tmp/dvc-storage directory and a metrics.json file will be produced with the evaluation metrics. These outputs are also cached by DVC so that subsequent executions can reuse them if the inputs haven‘t changed.

Now let‘s see how data version control helps with experiment tracking. Say we want to try a different set of hyperparameters for the Random Forest model. We can modify the params.yaml file:

train:
  n_estimators: 150 
  min_samples_split: 4

Then run the pipeline again:

dvc repro

DVC will detect that only the train and evaluate stages need to be rerun, since the outputs of the prepare stage haven‘t changed. It will execute those stages with the updated parameters, generate a new model.pkl file and metrics.json file.

To track this experiment, we can commit the updated params.yaml file and .dvc files to Git:

git add params.yaml model.pkl.dvc metrics.json.dvc
git commit -m "Experiment with increased n_estimators and min_samples_split"

This creates a new version in the Git history corresponding to this experiment. We can switch back to the previous version to compare the results:

git checkout HEAD~1
dvc checkout

The dvc checkout command pulls the previous data and model files from the DVC cache. We can then compare the metrics between the two experiments:

$ dvc metrics diff
Path          Metric    HEAD^  HEAD     Change
metrics.json  accuracy  0.967  0.983    0.017
metrics.json  f1        0.963  0.982    0.020

This shows that the increased n_estimators and min_samples_split hyperparameters improved the accuracy and F1 score by 1-2%.

Best Practices

Data version control is a powerful tool for MLOps, but like any tool it must be used properly to be effective. Here are some best practices to follow:

• Treat data as code: Commit changes to data in small, logical chunks just like you would commit code changes. Use meaningful commit messages to describe what changed.

• Use lightweight formats: DVC works well with simple formats like CSV or Parquet that can be split into smaller chunks. Avoid large monolithic files like relational databases that are hard to version.

• Adopt pipelines: Pipelines make your workflow reproducible by explicitly defining the stages in your ML lifecycle and the dependencies between them. Use a tool like DVC to create pipelines and track experiments.

• Track and compare metrics: Metrics are the key output of any ML experiment. Make sure to track relevant metrics in a machine readable format so you can easily compare them across experiments using Git tags or branching.

• Collaborate using pull requests: Data version control makes it easy to collaborate on ML experiments by creating branches or forks. Use pull requests to propose, discuss, and review changes before merging them into the main branch.

By following these practices, you can get the most value out of data version control and streamline your experiment workflow.

Conclusion

Data version control is an increasingly important part of the MLOps toolkit. By versioning data and models alongside code, it provides a single source of truth for ML experiments and makes them reproducible. Tools like DVC integrate with Git to track changes to data and models, and use pipelines to define reusable workflows.

In this post we walked through an example of using DVC to track experiments on the Iris dataset. We showed how to define a multistage pipeline, run experiments with different parameters, and compare the results. We also discussed best practices for making the most of data version control.

If you‘re getting started with MLOps, I recommend adding a data version control tool to your stack. Start by versioning a sample dataset and model and creating a simple pipeline. Over time you can expand to more complex use cases and workflows. By versioning your data and treating your models as code, you‘ll be able to collaborate more effectively and ensure your ML project is reproducible at every stage.