Saving and versioning models: a practical AI & ML guide

Introduction

You train a model that scores 94% accuracy on your validation set. Six months later, a teammate asks you to retrain it—but they can't reproduce your results because the training script has changed, dependencies have been upgraded, and nobody remembers which data snapshot you used. This is the cost of treating models like one-off experiments instead of build artifacts. Saving and versioning models: a practical AI & ML guide means we're applying the same rigor we use for code and infrastructure to the machine learning lifecycle.

In this post, we'll walk through serialization, versioning strategy, model registries, and reproducibility techniques that turn your models into reliable, auditable assets your team can depend on.

Serialize — joblib for scikit-learn, state_dict for PyTorch

Before you can version a model, you need to save it. The serialization format matters because it affects compatibility, file size, and what metadata you can attach.

scikit-learn and joblib

For scikit-learn pipelines, joblib is the standard choice. It handles pickle-based serialization and supports compression, which keeps file sizes manageable.

import joblib
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier

# After training your pipeline
pipeline = Pipeline([
    ('scaler', StandardScaler()),
    ('classifier', RandomForestClassifier(n_estimators=100, random_state=42))
])

# Train on your data
pipeline.fit(X_train, y_train)

# Save with compression
joblib.dump(pipeline, 'model_v1.pkl', compress=3)

# Load later
loaded_pipeline = joblib.load('model_v1.pkl')

The compress=3 flag reduces file size significantly—critical when you're managing hundreds of models. Joblib also handles most scikit-learn objects cleanly, including custom transformers if they follow the sklearn API.

PyTorch and state dictionaries

PyTorch models require a different approach. Save only the state_dict (weights and biases), not the entire model object, to avoid breaking changes when code evolves.

import torch
import torch.nn as nn

class SimpleClassifier(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        return self.fc2(x)

# After training
model = SimpleClassifier(input_size=10, hidden_size=64, output_size=3)
# ... training loop ...

# Save state dict, not the model object
torch.save(model.state_dict(), 'model_weights_v1.pt')

# Load and rebuild
model = SimpleClassifier(input_size=10, hidden_size=64, output_size=3)
model.load_state_dict(torch.load('model_weights_v1.pt'))
model.eval()

By saving the state dict separately from the model class definition, you keep flexibility. Your model code can change (fix a bug, refactor), and you can still load old weights by instantiating the class and loading the dict.

ONNX for cross-framework portability

If your team uses different frameworks—PyTorch for training, TensorFlow for inference—consider ONNX (Open Neural Network Exchange). It's a serialization format that most frameworks can read and write, reducing vendor lock-in.

Version everything — Model, code, data snapshot and metrics together

A model file alone is useless without context. You need to track the code that trained it, the data it saw, and the metrics it achieved. We call this bundled information a model artifact.

The versioning bundle

When you save a model, create a corresponding metadata file or directory structure that captures:

• Model weights: The serialized model file.
• Training code: Git commit hash or the exact script used.
• Data snapshot: A hash or pointer to the training dataset version.
• Hyperparameters: All settings used (learning rate, batch size, regularization).
• Metrics: Validation accuracy, precision, recall, AUC, loss curve.
• Dependencies: Versions of scikit-learn, torch, numpy, etc.
• Training timestamp: When the training completed.

Here's a practical approach: store metadata as JSON alongside your model file.

import json
import joblib
from datetime import datetime
import subprocess

def save_model_with_metadata(model, model_name, version, metrics, hyperparams):
    """
    Save model and metadata together.
    """
    # Get git commit hash
    git_hash = subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode().strip()
    
    metadata = {
        'model_name': model_name,
        'version': version,
        'git_commit': git_hash,
        'training_timestamp': datetime.utcnow().isoformat(),
        'metrics': metrics,
        'hyperparameters': hyperparams,
        'dependencies': {
            'scikit-learn': '1.3.2',
            'numpy': '1.24.3',
            'pandas': '2.0.3'
        }
    }
    
    # Save model
    model_path = f'{model_name}_v{version}.pkl'
    joblib.dump(model, model_path, compress=3)
    
    # Save metadata
    metadata_path = f'{model_name}_v{version}.json'
    with open(metadata_path, 'w') as f:
        json.dump(metadata, f, indent=2)
    
    return model_path, metadata_path

# Usage
metrics = {'val_accuracy': 0.94, 'val_precision': 0.92, 'val_recall': 0.91}
hyperparams = {'n_estimators': 100, 'max_depth': 15, 'random_state': 42}

save_model_with_metadata(
    pipeline, 
    'customer_churn', 
    version='1.0', 
    metrics=metrics, 
    hyperparams=hyperparams
)

This approach is lightweight and works well for smaller teams. For each model version, you get a pair of files that completely describe how it was trained and how it performed.

Model registry — MLflow tracks versions and stages

As you scale to dozens or hundreds of models, file-based versioning becomes tedious. A model registry centralizes storage, versioning, and lifecycle management. MLflow is the industry standard for this.

Setting up MLflow

MLflow provides a server and API for tracking experiments and registering models. Here's a minimal setup:

import mlflow
from mlflow.models import infer_signature
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

# Start an MLflow run
mlflow.start_run()

# Log hyperparameters
mlflow.log_param('n_estimators', 100)
mlflow.log_param('max_depth', 15)
mlflow.log_param('random_state', 42)

# Load data and train
X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

model = RandomForestClassifier(n_estimators=100, max_depth=15, random_state=42)
model.fit(X_train, y_train)

# Log metrics
accuracy = model.score(X_test, y_test)
mlflow.log_metric('test_accuracy', accuracy)

# Infer and log model signature
signature = infer_signature(X_test, model.predict(X_test))

# Register the model in MLflow registry
mlflow.sklearn.log_model(
    model,
    'iris_classifier',
    signature=signature,
    registered_model_name='iris_classifier'
)

mlflow.end_run()

Model stages and transitions

MLflow allows you to tag model versions with stages: None (development), Staging, Production, and Archived. This mirrors your deployment pipeline.

from mlflow.tracking import MlflowClient

client = MlflowClient()

# Transition version 2 to Staging for evaluation
client.transition_model_version_stage(
    name='iris_classifier',
    version=2,
    stage='Staging'
)

# After testing, move to Production
client.transition_model_version_stage(
    name='iris_classifier',
    version=2,
    stage='Production'
)

# Archive old production version
client.transition_model_version_stage(
    name='iris_classifier',
    version=1,
    stage='Archived'
)

This gives your team a single source of truth. Data scientists train models and log them to MLflow; the production team sees staging candidates and can safely promote them. Nobody manually copies model files around.

Reproducibility — Fix random seeds and pin library versions

Serialization and versioning are only useful if others can reproduce your results. Randomness is the enemy of reproducibility.

Set random seeds everywhere

Different libraries use different random number generators. Set them all:

import os
import numpy as np
import torch
import random
from sklearn.utils.random import check_random_state

def set_seed(seed=42):
    """Set all random seeds for reproducibility."""
    os.environ['PYTHONHASHSEED'] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# Call at the start of your training script
set_seed(42)

This ensures that every time you run the training script with the same data and hyperparameters, you get the same weights. Critical for debugging and for enabling teammates to verify results.

Pin dependency versions

A requirements.txt should be as strict as possible. Use exact versions, not ranges:

scikit-learn==1.3.2
numpy==1.24.3
pandas==2.0.3
torch==2.0.1
mlflow==2.7.0
joblib==1.3.2

For Docker deployments, this becomes even more critical. Your production inference container should run the exact same dependency versions as your training environment. Many model bugs in production are actually compatibility issues between library versions.

Document the data snapshot

Include a data version in your model metadata. This can be a Git hash (if your data is versioned), an S3 bucket key with a timestamp, or a database query hash:

import hashlib
import pandas as pd

def hash_dataframe(df):
    """Create a deterministic hash of a dataframe."""
    return hashlib.sha256(pd.util.hash_pandas_object(df, index=True).values).hexdigest()

# After loading training data
data_hash = hash_dataframe(X_train)
mlflow.log_param('training_data_hash', data_hash)

If a teammate wants to rebuild your model from scratch, they know exactly which data was used. If results drift over time, you can compare the data hash to see if the training distribution has changed.

Common mistakes

• Saving the entire model object instead of state_dict (PyTorch): When your model code changes, you can't load old weights. Always save state_dict.
• Skipping metadata: The model file means nothing without context. Log hyperparameters, metrics, data hash, and commit hash every time.
• Not pinning dependencies: If a library updates and breaks behavior, your training becomes irreproducible. Use exact versions in requirements.txt.
• Forgetting to set random seeds: Running the same training code twice produces different models. Your team loses confidence in reproducibility.
• Mixing model registry with file storage: Pick one approach and stick with it. MLflow is cleaner for teams; JSON metadata files work for small projects. Don't use both.
• No data versioning strategy: If you don't track which data snapshot trained each model, you can't trace bugs or performance regressions.

Key takeaways

• Use joblib for scikit-learn, PyTorch state_dict for neural networks. Serialization format matters.
• Bundle models with metadata: code commit, data hash, hyperparameters, and metrics. Version the whole package, not just weights.
• Adopt MLflow or a similar registry for teams. Stages (Staging, Production) let you safely promote models without manual file management.
• Set all random seeds in your training script. Pin dependency versions in requirements.txt. Reproducibility requires discipline.
• Treat models like production code: version control, code review, auditable history. Your future self will thank you.