Model registry in ML Python

Introduction

A model registry helps you keep track of different versions of your machine learning models in one place.

When you want to save and organize models during development.

When you need to share models with your team easily.

When you want to deploy the best model version to production.

When you want to compare model versions and their performance.

When you want to keep a history of model changes for future reference.

Syntax

ML Python

register_model(model, name, version=None, description=None)

# model: the trained ML model object
# name: a string to identify the model
# version: optional version number or string
# description: optional text about the model

The register_model function saves your model with a name and optional version.

Descriptions help explain what the model does or its changes.

Examples

Save the model with the name 'house_price_predictor' without specifying a version.

ML Python

register_model(my_model, 'house_price_predictor')

Save the model with a version label 'v1' to track this specific version.

ML Python

register_model(my_model, 'house_price_predictor', version='v1')

Add a description to explain what data or changes this model has.

ML Python

register_model(my_model, 'house_price_predictor', description='Model trained on 2024 data')

Sample Model

This code trains a simple linear regression model, saves it in a basic model registry, then loads it back to make predictions and calculate error.

ML Python

from sklearn.linear_model import LinearRegression
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

# Simple model registry as a dictionary
model_registry = {}

def register_model(model, name, version=None, description=None):
    key = f"{name}:{version}" if version else name
    model_registry[key] = {'model': model, 'description': description}
    print(f"Model '{key}' registered.")

def get_model(name, version=None):
    key = f"{name}:{version}" if version else name
    return model_registry.get(key, {}).get('model')

# Create sample data
X, y = make_regression(n_samples=100, n_features=1, noise=10, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)

# Train a simple linear regression model
model = LinearRegression()
model.fit(X_train, y_train)

# Register the model
register_model(model, 'simple_linear', version='v1', description='Linear regression on synthetic data')

# Retrieve the model and predict
loaded_model = get_model('simple_linear', version='v1')
predictions = loaded_model.predict(X_test)

# Calculate and print mean squared error
mse = mean_squared_error(y_test, predictions)
print(f"Mean Squared Error: {mse:.2f}")

Important Notes

A model registry can be a simple dictionary or a full system like MLflow or SageMaker.

Always include versioning to avoid confusion between models.

Descriptions help your team understand model differences quickly.

Summary

A model registry stores and organizes your machine learning models.

It helps track versions, share models, and deploy the best one.

Even a simple dictionary can work as a basic model registry.