Module 09 Intermediate 18 min read

Model Persistence

Save and load models with Pickle and Joblib, persist Keras models, version artefacts with metadata, and avoid common serialisation pitfalls.

Updated 2025 · Edit on GitHub

Why Model Persistence?

Training takes time and compute. Once trained, you need to save the model so it can be loaded later for prediction — without retraining. You also need to save all the preprocessors (scalers, encoders) that were fit on training data, so they can be applied identically to new inputs.

⚠️
Always save your preprocessors alongside your model. If you save the model but not the StandardScaler, you can't make predictions on new data — you'll apply raw features to a model trained on scaled ones. Package them together in a Pipeline.

Pickle — General Python Serialisation

Python
import pickle
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split

# Train
X, y = load_breast_cancer(return_X_y=True)
X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.2, random_state=42)

pipe = Pipeline([
    ("scaler", StandardScaler()),
    ("model",  GradientBoostingClassifier(n_estimators=100, random_state=42))
])
pipe.fit(X_tr, y_tr)
print(f"Test accuracy: {pipe.score(X_te, y_te):.4f}")

# ── Save
with open("models/model_v1.pkl", "wb") as f:
    pickle.dump(pipe, f, protocol=pickle.HIGHEST_PROTOCOL)
print("Model saved to models/model_v1.pkl")

# ── Load
with open("models/model_v1.pkl", "rb") as f:
    loaded_pipe = pickle.load(f)

# Identical predictions
import numpy as np
assert np.allclose(pipe.predict_proba(X_te), loaded_pipe.predict_proba(X_te))
print("Loaded model predictions match original ✓")
⚠️
Pickle security risk: Never unpickle files from untrusted sources — pickle can execute arbitrary code during deserialization. For production APIs receiving models from external sources, use joblib + signature verification or ONNX format.

Joblib — Optimised for NumPy Arrays

Joblib is the preferred tool for scikit-learn models. It uses a more efficient serialisation for large NumPy arrays (like trained tree ensembles) and supports compression out of the box:

Python
import joblib
import os

os.makedirs("models", exist_ok=True)

# ── Save with compression (compress=3 gives good speed/size balance)
joblib.dump(pipe, "models/model_v1.joblib", compress=3)
size_kb = os.path.getsize("models/model_v1.joblib") / 1024
print(f"Saved: {size_kb:.1f} KB")

# ── Load
loaded = joblib.load("models/model_v1.joblib")
print(f"Loaded model accuracy: {loaded.score(X_te, y_te):.4f}")

# ── Benchmark: Pickle vs Joblib
import time
for name, save_fn, load_fn, path in [
    ("Pickle", lambda p: pickle.dump(pipe, open(p,"wb")), lambda p: pickle.load(open(p,"rb")), "models/pkl.pkl"),
    ("Joblib", lambda p: joblib.dump(pipe, p, compress=3), lambda p: joblib.load(p), "models/jbl.joblib"),
]:
    t = time.perf_counter(); save_fn(path); save_t = time.perf_counter()-t
    t = time.perf_counter(); load_fn(path); load_t = time.perf_counter()-t
    size = os.path.getsize(path)/1024
    print(f"{name:8s} save:{save_t:.3f}s  load:{load_t:.3f}s  size:{size:.1f}KB")

Saving Keras / TensorFlow Models

Python
import tensorflow as tf
from tensorflow import keras

# Build and train a small model (example)
model = keras.Sequential([
    keras.layers.Dense(64, activation="relu", input_shape=(30,)),
    keras.layers.Dense(1,  activation="sigmoid")
])
model.compile(optimizer="adam", loss="binary_crossentropy")
model.fit(X_tr, y_tr, epochs=5, verbose=0)

# ── Recommended: SavedModel format (full model + weights + graph)
model.save("models/keras_model.keras")     # .keras format (TF 2.12+)
reloaded = keras.models.load_model("models/keras_model.keras")

# ── Weights only (need to rebuild architecture first)
model.save_weights("models/weights.h5")
model.load_weights("models/weights.h5")

# ── TF-Lite: compress for mobile/edge deployment
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
with open("models/model.tflite", "wb") as f:
    f.write(tflite_model)
print(f"TFLite size: {len(tflite_model)/1024:.1f} KB")

Model Versioning Best Practices

Python
import json
from datetime import datetime

# Save model with metadata
metadata = {
    "model_name":     "GradientBoosting-BreastCancer",
    "version":        "1.0.0",
    "trained_at":     datetime.now().isoformat(),
    "train_accuracy": float(pipe.score(X_tr, y_tr)),
    "test_accuracy":  float(pipe.score(X_te, y_te)),
    "n_features":     int(X.shape[1]),
    "sklearn_version": __import__("sklearn").__version__,
    "python_version":  __import__("sys").version,
}

os.makedirs("models/v1", exist_ok=True)
joblib.dump(pipe, "models/v1/pipeline.joblib")
with open("models/v1/metadata.json", "w") as f:
    json.dump(metadata, f, indent=2)

print("Saved model + metadata:")
print(json.dumps(metadata, indent=2))

Summary

  • Always save the full Pipeline (preprocessor + model together) — not just the model.
  • Pickle: general purpose. Joblib: better for NumPy-heavy models. .keras: TensorFlow models.
  • Save metadata (version, accuracy, feature count, library versions) alongside every model artefact.
  • Never unpickle files from untrusted sources. Consider ONNX for cross-framework portability.