Tracking Your Machine Learning Project Changes with Neptune
Machine learning (ML) powers intelligent applications across industries, from retail recommendation engines to autonomous vehicles. But the process of developing and deploying ML models is far from easy. Data scientists and ML engineers face numerous challenges including:
- Managing experiments across many iterations, hyperparameters, and datasets
- Debugging errors and unexpected model behavior
- Reproducing and explaining results to collaborate with team members
- Auditing models to ensure compliance with data regulations
- Deploying models to production and monitoring performance over time
Disciplined experiment tracking is critical to overcoming these challenges. By diligently logging the data, code, environment, hyperparameters, and evaluation metrics for each experiment, ML teams can more efficiently iterate on models and ensure smooth handoffs from development to production.
Neptune is a metadata store designed for tracking and organizing ML experiments. In this post, we‘ll walk through how to use Neptune to track changes across your model development workflow.
The growing challenge of ML experiment tracking
Developing a model to solve a supervised learning task on a static dataset is fairly straightforward. The real challenge is iterating and scaling that model to meet the changing needs of the business. This is where experiment tracking becomes essential.
In a 2021 survey of ML practitioners, 84% reported that experiment tracking is critical for their work. However, over 40% still rely on spreadsheets or text files to manually log their work.
Experiment tracking becomes exponentially harder as the ML workflow grows in complexity:
- More people collaborating on model development
- Larger datasets that change over time
- New model architectures and hyperparameters to tune
- Continuous integration and deployment (CI/CD) pipelines
- Multiple target environments (cloud, edge, mobile)
- More stringent model auditing and compliance requirements
Without a system to track all the moving pieces, ML teams risk wasting time on duplicate work, struggling to reproduce results, and making costly model debugging or deployment mistakes.
The 2021 Algorithmia survey of ML teams found that 64% of respondents take a month or longer to deploy a model, with experiment tracking issues contributing to those delays. And a 2019 Dotscience survey estimates that data scientists spend up to 40% of their time on "bad" or unproductive data science, with 87% reporting negative business impacts as a result.
The downstream impacts of poor experiment tracking are clear. That‘s why thoughtful tracking is a key component of MLOps, the discipline of collaboration between data science and IT teams to deploy and maintain ML systems in production reliably and efficiently.
Tracking experiments with Neptune
Neptune is a leading tool for experiment management, metadata versioning, model registry, and MLOps monitoring. Users can instrument their ML runs to log over 25 types of metadata to Neptune, including:
- Parameters
- Metrics
- Code and notebook versions
- Data versions
- Model binaries
- Interactive visualizations
- Audio
- Video
Neptune‘s flexible metadata store scales with your ML workflow maturity. Individual data scientists can use it to track ad hoc experiments on their laptop. As the work scales to production ML systems, Neptune tracks the full lineage of experiments across distributed compute infrastructure all the way to model monitoring in production.
Some of the key benefits of using Neptune include:
- Experiment Management: Visibility into all ML experiments in a central UI to search, group, and compare across model versions.
- Model Registry: Persistent model store to version, describe, and share models across your organization.
- ML Monitoring: Catch model performance degradation, data drift, and latency issues in production.
- Team Collaboration: Share experiment results and reports across data science and business teams. Extensive access controls and audit logging.
- Workflow Integration: Neptune plugs into your existing ML workflows via flexible APIs and integrations with frameworks like PyTorch, Tensorflow, Keras, Scikit-Learn, and more.
- Flexible Deployment: Neptune can run as a fully-managed SaaS or self-hosted in your own cloud or on-prem infrastructure.
Now let‘s get our hands dirty and see how to implement Neptune in a realistic ML project.
A tour of ML experiment tracking in Neptune
We‘ll use the classic UCI Iris flower dataset to build a RandomForestClassifier model predicting flower species from petal and sepal measurements.
Installing Neptune
First, install the Neptune client library in your Python environment:
pip install neptune-clientNext, create a free Neptune AI account. Add your Neptune API token to an environment variable for easy access:
export NEPTUNE_API_TOKEN=<YOUR_API_TOKEN>Creating a Neptune run
In your Python script or notebook, create a Neptune Run object to track metadata for a single experiment:
import neptune.new as neptune
run = neptune.init(
project=‘YourUserName/iris-classification‘,
api_token=os.getenv(‘NEPTUNE_API_TOKEN‘),
)The project argument takes the form of <WORKSPACE/PROJECT>, where WORKSPACE is your username or organization name, and PROJECT is the name of your ML project.
Logging run metadata
Now we can log information about the ML run to Neptune. Let‘s track the ModelCard metadata about the dataset:
from sklearn.datasets import load_iris
data = load_iris(as_frame=True)
X, y = data.data, data.target
run[‘dataset/name‘] = data.frame.dataset_name
run[‘dataset/num_classes‘] = len(data.frame.species.unique())
run[‘dataset/size‘] = len(X)
run[‘dataset/features‘] = list(X.columns)
run[‘dataset/version‘] = 1 We log model hyperparameters in a similar fashion:
from sklearn.ensemble import RandomForestClassifier
params = {
‘n_estimators‘: 100,
‘max_depth‘: 5,
‘min_samples_split‘: 2,
‘min_samples_leaf‘: 1,
}
clf = RandomForestClassifier(**params)
run[‘model/params‘] = params
run[‘model/name‘] = clf.__class__.__name__As we train and evaluate the model, we can log metrics for each run:
from sklearn.model_selection import cross_validate
cv_results = cross_validate(clf, X, y, scoring=[‘f1_macro‘, ‘precision_macro‘, ‘recall_macro‘])
run[‘metrics/train/f1_macro‘] = cv_results[‘test_f1_macro‘].mean()
run[‘metrics/train/precision‘] = cv_results[‘test_precision_macro‘].mean()
run[‘metrics/train/recall‘] = cv_results[‘test_recall_macro‘].mean()Finally, we can log the trained model binary for later use:
import joblib
joblib.dump(clf, ‘model.pkl‘)
run[‘artifacts/model‘].upload(‘model.pkl‘)Querying and comparing runs
After tracking a few experiments, we can use the Neptune UI to browse our runs.
Neptune‘s table view lets us quickly compare runs across hyperparameters and metrics. We can filter, group, and sort to find our best performing models.
Clicking on a run lets us see more details like the exact code version, dataset samples, model binary, and hardware environment. This information is crucial for reproducing and debugging results.
Visualizations and reporting
In addition to logging model inputs and outputs, we can log custom visualizations to Neptune to better understand our results.
For example, here‘s how to log a confusion matrix for our Iris classifier using Scikit-Learn and Matplotlib:
from scikitplot.metrics import plot_confusion_matrix
import matplotlib.pyplot as plt
_, ax = plt.subplots(figsize=(12, 12))
plot_confusion_matrix(y_test, y_pred, ax=ax)
run[‘plots/confusion_matrix‘].upload(File.as_image(fig))
These visualizations are a compelling way to share ML results with collaborators and business stakeholders. We can combine them with text and Neptune runs to build dashboards and reports right within the Neptune UI.
Advanced logging
So far we‘ve covered basic Neptune logging functionality. Where Neptune really shines is in tracking more complex model and data workflows.
For example, here‘s how to track your data preprocessing code and resulting dataset versions:
import pandas as pd
data = pd.read_csv(‘data.csv‘)
X = preprocess(data)
run["dataset/v2/preprocessed"].upload(File.as_html(X))
run["dataset/v2/code/preprocessing.py"].upload(‘preprocessing.py‘)By versioning our datasets with the accompanying processing code, we can ensure full reproducibility and audit how data evolves over the course of the project.
Here‘s an example of logging an entire model training pipeline built with Scikit-Learn:
from sklearn.pipeline import Pipeline
pipe = Pipeline([
(‘vectorizer‘, TfidfVectorizer()),
(‘classifier‘, RandomForestClassifier())
])
run[‘model/pipeline‘] = pipe
run[‘model/pipeline/params‘] = pipe.get_params()
pipe.fit(text_data, labels)
run[‘model/pipeline/trained‘] = pipeLogging the pipeline as a single artifact provides a complete picture of the model, making it easier to move to production.
Exploring the Neptune ecosystem
The core Neptune library provides a simple, flexible API for logging all types of metadata. On top of this foundation, Neptune offers numerous integrations to fit your existing ML tools and workflow:
- Neptune-PyTorch and Neptune-TensorFlow integrations make it easy to track deep learning model training with advanced features like learning curves and GPU utilization.
- Neptune-MLflow integration syncs MLflow runs to Neptune for richer tracking and visualization capabilities.
- R and Julia SDKs lets users in R and Julia log metadata to Neptune.
- NeptuneContrib is a collection of powerful Neptune utilities like hyperparameter sweeps and interactive plots.
The Neptune integrations page lists dozens more tools in the ML ecosystem that Neptune connects with.
Conclusion
Experiment tracking is a crucial part of the ML workflow. Without it, data science teams waste time and lack reproducibility.
In this post, we saw how Neptune makes it easy to track all types of ML metadata in a single place. By logging parameters, metrics, visualizations, dataset versions, and more to Neptune, we gain full visibility into each model experiment. That translates to more productive iteration, easier debugging, and smoother collaboration.
Getting started with Neptune is easy and free. Install the Python library, create a Neptune account, and start logging your ML experiments. As you mature your workflows with more advanced technology, Neptune will scale with you on the journey to production ML and automated MLOps.
Further Reading
To learn more about ML experiment tracking with Neptune, check out these resources:
