Best ML Model Management Tools

Working with Machine Learning models requires a different set of organizational capabilities as opposed to building software. These new sets of capabilities are driven by MLOps – which is further divided into different ML model processes, one of which is Model Management. 

Machine Learning model management is a class of technologies and processes that work together to enable companies to reliably and securely develop, validate, deliver, and monitor models at high velocity, mainly to create a competitive advantage. It covers almost the entire spectrum of managing models from development to re-training and monitoring. As a result, model management can be divided into a few phases:

• Optimization 
• Versioning and tracking
• Evaluation and testing
• Packaging
• Deployment and serving
• Monitoring

ML model optimization

Hyperparameter optimization is a key step in model management. To get the best performance out of your model, you have to tweak the hyperparameters to get the most out of the data. There are several tools available that can help you achieve this goal. Let’s explore a few options that tick most of the boxes:

1. Optuna

Optuna is an automatic hyperparameter optimization framework, particularly designed for machine learning. It features an imperative, define-by-run style user API hence it’s highly modular, and the coding functionality can also be used in other domains.

Optuna features:

• Lightweight, versatile, and platform-agnostic architecture: Handle a wide variety of tasks with a simple installation that has few requirements.
• Pythonic search spaces: Define search spaces using familiar Python syntax, including conditionals and loops.
• Efficient optimization algorithms: Adopt state-of-the-art algorithms for sampling hyperparameters and efficiently pruning unpromising trials.
• Easy parallelization: Scale studies to tens or hundreds of workers with little or no changes to the code.
• Quick visualization: Inspect optimization histories from a variety of plotting functions.

2. Sigopt

SigOpt is a model development platform that makes it easy to track runs, visualize training, and scale hyperparameter optimization for any type of model built with any library on any infrastructure. 

SigOpt optimizes models (and even advertising and products) through a combination of Bayesian and global optimization algorithms, which allows users and companies to boost the performance of their models while cutting costs and saving the time that might have been previously spent on more tedious, conventional tuning methods.

SigOpt features:

• Robust tracking: SigOpt tracks and organizes your training and tuning cycles, including architectures, metrics, parameters, hyperparameters, code snapshots, and the results of feature analysis, training runs, or tuning jobs.
• Automated Training & Tuning: It fully integrates automated hyperparameter tuning with training run tracking to make this process easy and accessible. And features like automated early stopping, highly customizable search spaces, and multimetric and multitask optimization make tuning easy for any model you are building.
• Easy Integrations: Sigopt offers easy integrations with all well-known frameworks like H2O, Tensorflow, PyTorch, scikit, etc.

Dig deeper

If you want to learn more, check our article about best tools for model tuning and hyperparameter optimization.

Versioning and tracking ML models

Amongst innumerable experiments, how do we keep track of changes in the data, source code, and ML models? Versioning tools offer a solution to this and are critical to the workflow as they cater to reproducibility. They help you get a past version of an artefact, dataset hash, or model that can be used to roll back for performance reasons or for the sake of comparison. 

You’d often log this data and model version into your metadata management solution to ensure your model training is versioned and reproducible. A versioned model will then have to be tracked in order to compare results from other versioned experiments.

In an ideal ML workflow, you’d want to know about how the metadata of your model was built, which ideas were tested, or where you can find all the packaged models. To achieve that, you need to keep track of all your model-building metadata like hyperparameters, metrics, code and data versions, evaluation results, packaged models, and other information.

Here are a few tools worth checking out.

1. Neptune

Neptune is an ML metadata store used for experiment tracking and model registry. You can log and display pretty much any ML metadata, from source code and artifacts to metrics, hyperparameters, model checkpoints, and system metrics. Having everything tracked in Neptune, you’re able to go downstream: from the model, through the run, to the dataset it was trained on, whenever it’s necessary.

The web UI was built for managing ML model experiments and models and lets you:

• search experiments and models by anything you want with an advanced query language
• customize which metadata you see with flexible table views and dashboards
• monitor, visualize, and compare experiments.

Besides logging, you can also query metadata from Neptune.

If you are wondering if it will fit your workflow, check out case studies of how people set up their MLOps tool stack with Neptune.

Neptune—summary:

• Flexible folder-like metadata structure
• Highly customizable UI for searching, displaying, comparing, and organizing ML metadata
• Easy-to-use API for logging and querying model metadata and 25+ integrations with MLOps tools
• Both hosted, and on-prem versions are available
• Collaboration features and organization/project/user management

2. DVC

Data Version Control(DVC) is essentially an experiment management tool for ML projects. DVC software is built upon Git, which means it lets you capture the versions of your data and models in Git commits while storing them on-premises or in cloud storage. and its main goal is to codify data, models, and pipelines through the command line. 

A few key features of DVC:

1. Lightweight: DVC is a free, open-source command line tool that doesn’t require databases, servers, or any other special services.
2. Consistency: Keep your projects readable with stable file names — they don’t need to change because they represent variable data. There is no need for complicated paths or constantly editing them in the source code.
3. Efficient data management: DVC stores big amounts of data in a well-organized and accessible way. This tremendously helps with Data versioning.
4. Collaboration: Easily distribute your project development and share its data internally and remotely, or reuse it in other places. 
5. Data compliance: Review data modification attempts as Git pull requests. Audit the project’s immutable history to learn when datasets or models were approved and why.
6. GitOps: Connect your data science projects with the Git ecosystem. Git workflows open the door to advanced CI/CD tools (like CML), specialized patterns such as data registries, and other best practices.
7. Bonus stuff: DVC offers some experiment tracking capabilities as well. It can track these experiments, list and compare their most relevant metrics, parameters, and dependencies, navigate among them and commit only the ones that are needed to Git.

3. Pachyderm

Pachyderm offers data versioning and pipelines for MLOps. They provide a data foundation that allows data science teams to automate and scale their machine learning lifecycle while guaranteeing reproducibility. It offers a commercial Enterprise Edition and an open-source Community Edition. If required, Pachyderm can be integrated into your existing infrastructure or private cloud as well.

Key features that Pachyderm aims to deliver are:

1. Cost-Effective Scalability: Delivering reliable results by optimizing resource utilization and running complex data pipelines with smart data transformation using auto-scaling and parallelism.
2. Reproducibility: Ensuring reproducibility and compliance via immutable data lineage and data versioning for any type of data. Git-like commits will increase team efficiency.
3. Flexibility: Pachyderm can leverage existing infrastructure and can run on your existing cloud or on-premises infrastructure. 

With the enterprise version, you get some additional features on top of the ones mentioned above:

• Authentication: Pachyderm allows for authentication against any OIDC provider. Users can authenticate to Pachyderm by logging into their favorite Identity Provider.
• Role-Based Access Control (RBAC): Enterprise-scale deployments require access control. Pachyderm Enterprise Edition gives teams the ability to control access to production pipelines and data. Administrators can silo data, prevent unintended modifications to production pipelines, and support multiple data scientists or even multiple data science groups by controlling users’ access to Pachyderm resources.
• Enterprise Server: An organization can have many Pachyderm clusters registered with one single Enterprise Server that manages the enterprise licensing and the integration with a company’s Identity Provider.
• Additionally, you have access to a pachctl command that pauses (pachctl enterprise pause) and unpauses (pachctl enterprise unpause) your cluster for a backup and restore.

4. MLflow

MLflow is a platform to streamline machine learning development, including tracking experiments, packaging code into reproducible runs, and sharing and deploying models. MLflow offers a set of lightweight APIs that can be used with any existing machine learning application or library (TensorFlow, PyTorch, XGBoost, etc.), wherever you currently run ML code (e.g., in notebooks, standalone applications, or the cloud). 

MLflow is available both as an open-source as well as a managed version (by Databricks). You can see a detailed comparison of the offerings of both versions here to make an informed decision.

The tool is library-agnostic. You can use it with any machine learning library or any programming language. MLflow comprises four main functions that help to track and organize experiments:

1. MLflow Tracking – Automatically logs parameters, code versions, metrics, and artifacts for each run using Python, REST, R API, and Java API
2. MLflow Projects – Represents a standard format for packaging reusable data science code.
3. MLflow Models – A standard format for packaging machine learning models that can be used in a variety of downstream tools—for example, real-time serving through a REST API or batch inference on Apache Spark.
4. MLflow Model Registry – offers a centralized model store, set of APIs, and UI, to collaboratively manage the full lifecycle of an MLflow Model.

Dig deeper

Check out a deeper comparison of tools for data versioning or take a look at the best experiment tracking and management tools.

ML model evaluation and testing

Evaluation makes up an important phase in ML workflow. Evaluation in ML often starts with defining metrics such as accuracy and confusion matrix depending on your data and use case. This level of model evaluation gets more sophisticated for real-world applications to take into account any data/concept drift.

While most teams are comfortable with using model evaluation metrics to quantify a model’s performance before deploying it, these metrics are mostly not enough to ensure your models are ready for production in real-world scenarios. Thorough testing of your models is required to ensure they are robust enough for real-world encounters. 

Now, the question is what to look for in the sea of tools. Here are some techniques that you will need while testing your model for the real world:

1. Unit Tests: Focusing on a particular scenario or subclass, e.g. “cars from the rear in low lighting”.
2. Regression Tests: Improvements obtained by a particular push (squashed bug) or one segment of the “long tail”.
3. Integration Tests: Mirroring deployment conditions in terms of composition, distribution, preprocessing, etc.

Let’s look at some tools that offer such ML model testing capabilities:

1. Kolena.io

Kolena is a machine learning testing and debugging platform to surface hidden model behaviors and take the mystery out of model development. Kolena helps you:

• Perform high-resolution model evaluation
• Understand and track behavioral improvements and regressions
• Meaningfully communicate model capabilities

They offer Unit tests, regression tests, and Integration tests as part of the package. You can create custom workflows for your use cases, however, there are some templates available, like image segmentation, object detection, etc. 

Kolena features:

1. Offers a host of testing capabilities:
   1. Unit/Regression/Integration Testing
   2. Option for combining multiple tests into a Test Suite
2. Can be accessed through a pip package – kolena-client.
3. Web dashboards are available to quickly visualize results.

Check out their documentation here to explore more.

2. Deepchecks

Deepchecks is the leading tool for testing and validating machine learning models and data, and it enables doing so with minimal effort. It caters to various validation and testing needs, such as verifying your data’s integrity, inspecting its distributions, validating data splits, evaluating your model, and comparing different models.

They offer tests for different phases in the ML workflow:

• Data Integrity
• Train-Test Validation (Distribution and Methodology Checks)
• Model Performance Evaluation

Deepchecks is an open-source tool that offers:

1. Tests for tabular and image data.
2. Integrations to Spark, Databricks, Airflow, etc. Check the full list here.
3. Results can be visualized in the working notebook or a separate window. Visit the documentation page here to understand how this works.
4. Unit/Regression/Integration tests.

Packaging ML models

A logical next step towards productionizing your ML model is to bundle/package it, so it behaves like an application. Your model sitting in a juypter lab environment won’t be able to serve users; by packaging it along with its dependencies, you are taking the first step towards exposing your model to the real world.

There are essentially two ways you could package your ML models:

1. Web-based frameworks: include microframeworks like Flask and FastAPI.
2. Containerization: includes docker, containerd, podman, etc.

Let’s look into some of the packaging tools:

1. Flask

Flask is a web framework, essentially a Python module, that lets you develop web applications easily. It has a small and easy-to-extend core: it’s a microframework that doesn’t include an ORM (Object Relational Manager) or such features. Flask depends on the Werkzeug WSGI toolbox and Jinja2 template. The ultimate purpose is to help develop a strong web application base.

Flask key features:

• It has a built-in development server and a fast debugger. It has a RESTful for dispatching of requests and request handling for HTTP
• It has inherent unit testing support.
• Multiple extensions provided by the community ease the integration of new functionalities.

2. FastAPI

FastAPI is a modern web-based framework for API building with Python version 3.6+ based on standard Python types, which provides high performance. It is an ASGI framework which makes it highly concurrent.

Fast API key features:

• Provides good performance, close to NodeJS and Go, as it uses Starlette and Pydantic. Among the fastest Python frameworks present.
• Is a robust framework with amazing and interactive documentation.
• Fully compatible with open standards for APIs, OpenAPI, and JSON Schema.
• Intuitive – great editor support, auto-completion everywhere, less debugging time.

3. Docker

Docker is an open-source containerization platform with an enterprise tier. It enables developers to package applications into containers—standardized executable components combining application source code with the operating system (OS) libraries and dependencies required to run that code in any environment.

Docker key features – 

1. Docker images can be deployed on the cloud using AWS ECS, Google Kubernetes Engine, etc.
2. Docker comes with CLI as well as applications for Windows, Mac, and Linux.
3. Development integrations are available in all popular code editors such as VS Code, Jetbrains, etc.
4. Software integrations are also available with all major cloud service providers.

ML model deployment and serving

After packaging your models, the immediate next step is to identify requirements in deployment and serving. Productionizing machine learning models presents its own challenges around operationalization.

The first batch of problems you hit is how to package, deploy, serve and scale infrastructure to support models in production. The tools for model deployment and serving help with that.

Here are a few model deployment and serving tools.

1. AWS
2. Cortex
3. BentoML

1. Amazon Web Services

Amazon web services is a leading cloud service provider with a market share of 33%. It offers over 200 fully featured services from data centres, and infrastructure technologies like compute, storage, and databases around the globe to emerging technologies, such as machine learning and artificial intelligence, data lakes and analytics, and the Internet of Things. This makes it faster, easier, and more cost-effective to move your existing applications to the cloud and build nearly anything you can imagine.

AWS features – 

1. It’s a big tech platform providing a host of infrastructure technologies.
2. Being a leader in the space it also provides highly efficient services to cater to your ML model needs.
3. AWS Elastic container registry is compatible with docker containers and seamlessly integrates with other services like AWS ECS or Fargate to put the container into production.
4. AWS Lambda lets you deploy your web-service serverless so you don’t have to worry about increased/decreased traffic as it will handle it automatically, thus saving costs and effort.
5. AWS Sagemaker offers a suite of services targeted toward ML development, deployment, and serving.

2. Cortex

Cortex is an open-source platform designed for deploying trained machine learning models directly as a web service in production. It is built on top of Kubernetes to support large-scale machine learning workloads. 

Cortex features:

• Multi framework: Cortex supports multiple frameworks, including TensorFlow, PyTorch, scikit-learn, XGBoost, and more.
• Autoscaling: Automatic load balancing for APIs to handle production workloads.
• Infrastructure: Cortex can run inference both on CPU or GPU infrastructure.
• Rolling updates: Cortex updates APIs post-deployment without any downtime.

3. BentoML

BentoML is an end-to-end solution for machine learning model serving. It facilitates Data Science teams to develop production-ready models serving endpoints, with best MLOps practices and performance optimization at every stage.

It offers a flexible and performant framework to serve, manage, and deploy ML models in production. It simplifies the process of building production-ready model API endpoints through its standard and easy architecture. It empowers teams with a powerful dashboard to help organize models and monitor deployments centrally.

BentoML features:

• Native support for popular ML frameworks: Tensorflow, PyTorch, XGBoost, Scikit-Learn, and many more!
• Custom pipelines: Define custom serving pipeline with pre-processing, post-processing, and ensemble models.
• Easy Versioning and deployment: Standard .bento format for packaging code, models, and dependencies for easy versioning and deployment.
• State-of-the-art Model Serving: BentoML offers online serving via REST API or gRPC, offline scoring on batch datasets with Apache Spark, or Dask, and stream serving with Kafka, Beam, and Flink.

Dig deeper

If you want to learn more, read our deeper comparison of tools for model deployment/serving.

ML model monitoring

Monitoring is another integral part of model management. There are several notions as to why monitoring your ML model is so necessary. To list a few:

1. To check how your model is actually performing with the real-world data. Is it producing crazy and unexpected results?
2. To check for a possible model drift.
3. Is it able to solve the business problem it was developed for in the first place?

Model monitoring tools can give you visibility into what is happening in production by:

• Monitoring input data drift: is your model getting different inputs as compared to the training data?
• Monitoring concept drift: is the problem no longer similar to what it initially was, and your model performance is decaying?
• Monitoring model drift: have the model’s predictive abilities degraded due to changes in the environment?
• Monitoring hardware metrics: is the new model much more resource hungry than the older one? 

Here’s what you can choose from.

1. Amazon SageMaker 

AWS Sagemaker is part of the AWS platform. Data scientists and developers can prepare, build, train, and deploy high-quality machine learning (ML) models. It claims to be the first integrated development environment (IDE) for ML. It has four components: prepare, build, train & tune, deploy & manage. Users can log, organize and compare experiments, automate hyperparameter tuning, and debug training runs.

Amazon SageMaker comes with a host of features, some of which are:

• Distributed Training: Sagemaker makes it faster for its users to perform the distributed training by splitting the data across multiple GPUs which achieves near-linear scaling efficiency.
• Profiling and Debugging Training Runs: SageMaker Debugger captures the metrics and profiles training jobs in real-time enabling users to correct performance problems quickly before the model is deployed to production.
• Supports Reinforcement Learning: It has built-in, fully-managed reinforcement learning algorithms including the newest and best performing in academic literature.
• Optimized for Major Frameworks: SageMaker is optimized for many popular deep learning frameworks like TensorFlow, Apache MXNet, PyTorch etc. Frameworks supported are always up-to-date with the latest version and are further optimized for performance on AWS.
• Supports AutoML: SageMaker autopilot automatically builds, trains, and tunes the best machine learning models, based on the users’ data while allowing them to maintain full control and visibility. The model can be deployed to production with just one click or iterated to further improve the model quality.

2. Fiddler

Fiddler is a Model Performance Management platform that offers model monitoring, observability, explainability & fairness. It lets you build a centralized platform to monitor and analyze AI models with simple plug-in integration and helps businesses understand the ‘why’ and ‘how’ behind their AI. The platform monitors AI models for data drift, data integrity issues, outliers, and performance drops.

Fiddler ML key features:

• Performance monitoring: It offers performance and evaluation metrics out-of-the-box for machine learning problems, including but not limited to binary classification, multi-class classification, regression, and ranking models.
• Data drift: With Fiddler, you can easily monitor data drift, uncover data integrity issues, and compare data distributions between baseline and production datasets to boost model performance.
• NLP and CV monitoring: Fiddler offers the monitoring of complex and unstructured data, such as natural language processing and computer vision.
• Alerts: It offers real-time alerts to minimize issues caused by performance, data drift, data integrity, and service metrics such as traffic, latency, and errors.

3. Arize

Arize is the machine learning observability platform for ML practitioners to monitor, troubleshoot, and explain models. It logs model inferences across training, validation, and production environments.

Arize AI has the following features: 

• ML Performance Tracking: Arize AI offers ML performance tracking with purpose build workflows to efficiently track and trace models. 
• Unstructured Data Monitoring: It supports unstructured data monitoring for NLP and computer vision.
• Understand Drift Impact: Real-world data is dynamic and impacts model response over time. With Arize you can track distribution changes in upstream data, predictions and actuals to proactively gauge model performance and find retraining opportunities.
• Automated Model Monitoring: Arize offers the ability to monitor every aspect of an ML model which is critical to catching performance degradation of key metrics. 
• Easy Integration & Deployment: Arize is designed to seamlessly plug into your existing ML stack and alerting workflows. The platform works with any model framework, from any platform, in any environment.

4. WhyLabs

WhyLabs is an AI observability platform designed to prevent data quality or model performance degradation by allowing you to monitor your data pipelines and machine learning models in production. 

Built on top of an open-source package called whylogs, it monitors data pipelines and ML applications for data quality regressions, data drift, and model performance degradation.

Key features of WhyLabs tool:

• Easy integration: Whylabs offers an open source module called whylogs which lets you integrate your code easily with WhyLabs SaaS platform.
• Data health / DataOps: The tool catches data quality issues automatically. It identifies and monitors data drifts and data bias before any impact on the user experience. 
• Model health / ModelOps: You can track model outputs and model performance, debug model and identify anomalies with smart correlation and visualization. Whylabs also alerts about concept drift and model accuracy degradation.
• Preserving privacy: WhyLabs profiles model inputs and outputs to capture only statistical profiles of the underlining data. The raw data never leaves the customer VPC/perimeter. All WhyLabs product features operate on statistical profiles.
• Zero maintenance: Whylabs offer a minimal to no maintenance structure with no schema maintenance, no monitoring configuration, no data sampling overhead, and no deployment pain.
• Seamless integration with your existing pipelines and tools:  Whylabs support all major frameworks like Tensorflow, PyTorch, scikit, Apache Spark, Azure, SageMaker etc.

Dig deeper

For deeper dive into monitoring frameworks, read our comparison of model monitoring tools.

Tools that can do it all (almost)!

If you’re looking for an all-around contender for your Model Management needs, then some options work as a “one-stop shop”. Let’s look at some such tools that offer services for a complete lifecycle.

1. Amazon SageMaker

As discussed above, Amazon SageMaker is a fully managed machine learning service. With SageMaker, you can:

1. quickly and easily build and train machine learning models, and 
2. then directly deploy them into a production-ready hosted environment. 

Sagemaker provides: 

1. An integrated Jupyter authoring notebook instance for easy access to your data sources for exploration and analysis, so you don’t have to manage servers. 
2. It also provides common machine learning algorithms that are optimized to run efficiently against extremely large data in a distributed environment. With native support for bring-your-own-algorithms and frameworks, 
3. SageMaker offers flexible distributed training options that adjust to your specific workflows. Deploy a model into a secure and scalable environment by launching it with a few clicks from SageMaker Studio or the SageMaker console. Training and hosting are billed by minutes of usage, with no minimum fees and no upfront commitments.

Sagemaker has a long list of feature offerings that covers development, packaging, deployment, and monitoring. Some services include Sagemaker Studio, Sagemaker Model Registry, Sagemaker Feature store, SagemakerExperiments, and Sagemaker Preprocessing. You can look at the full features list here.

2. Azure ML

Azure Machine Learning is a cloud service for managing the machine learning project lifecycle. It’s an all-round tool to Train, Deploy models, and manage MLOps. You can either create a model in Azure Machine Learning or use a model built from an open-source platform, such as Pytorch, TensorFlow, or scikit-learn. 

Azure ML offers:

• Collaboration, such as shared notebooks, compute resources, data, and environments, tracking, and auditability that shows who made changes and when, and asset versioning.
• Azure Machine Learning studio, which is a graphical user interface for the project workspace allows you to:
  • View runs, metrics, logs, outputs, and so on.
  • Author and edit notebooks and files.
  • Manage common assets, such as data credentials, compute and environments.
  • Visualize run metrics, results, and reports.
  • Visualize pipelines authored through developer interfaces.
  • Author AutoML jobs.
• Extensive support for ML frameworks such as PyTorch, Tensorflow, scikit-learn, XGBoost, LightGBM etc.
• Integrations enabling MLOPs:
  • Git integration
  • MLflow integration
  • Machine learning pipeline scheduling
  • Azure Event Grid integration for custom triggers
  • Easy to use with CI/CD tools like GitHub Actions or Azure DevOps

Note: Azure ML studio will be integrated into Azure Machine Learning by 2024.

3. Vertex AI Google Cloud

Vertex AI is a managed machine learning platform that provides you with all of Google’s cloud services in one place to deploy and maintain ML models. It brings together the Google Cloud services for building ML under one unified UI and API. 

In Vertex AI, you can now easily train and compare models using AutoML or custom code training and all your models are stored in one central model repository. These models can now be deployed to the same endpoints on Vertex AI.

Vertex AI offers:

• Training and hyperparameter tuning: Vertex AI Training offers fully managed training services, and Vertex AI Vizier provides optimized hyperparameters for maximum predictive accuracy.
• Model serving: Vertex AI Prediction makes it easy to deploy models into production, for online serving via HTTP or batch prediction for bulk scoring.
• Model tuning and understanding: Vertex Explainable AI tells you how important each input feature is to your prediction.
• Model monitoring: Continuous monitoring offers easy and proactive monitoring of model performance over time for models deployed in the Vertex AI Prediction service.
• Model management: Vertex ML Metadata enables easier auditability and governance by automatically tracking inputs and outputs to all components in Vertex Pipelines for artefact, lineage, and execution tracking for your ML workflow.