Saving and Loading Models

Persist a trained model to disk so you can reuse it without retraining.

Syntaxjoblib.dump(model, 'model.joblib')

Training can be slow, so you save the fitted model and load it later to make predictions — the basis of putting a model into production.

Tools

joblib is preferred for scikit-learn models (efficient with NumPy arrays). The standard-library pickle also works.

Example

Example · python
import joblib
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris

X, y = load_iris(return_X_y=True)
model = RandomForestClassifier().fit(X, y)

joblib.dump(model, 'model.joblib')     # save

loaded = joblib.load('model.joblib')   # load later
print(loaded.predict(X[:1]))           # [0]

When to use it

  • An ML engineer saves a trained scikit-learn pipeline with joblib after the overnight training run, then loads it in the Flask API without retraining.
  • A deep learning researcher saves PyTorch model weights as a .pt file after each epoch's validation improvement to keep the best checkpoint.
  • A data scientist versions saved model files by embedding the date and accuracy score in the filename so past models can be compared and reloaded.

More examples

Save and load with joblib

Serialises a trained scikit-learn model to disk with joblib.dump and confirms it predicts identically after being reloaded with joblib.load.

Example · python
import joblib
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris

X, y = load_iris(return_X_y=True)
model = RandomForestClassifier(n_estimators=50).fit(X, y)

joblib.dump(model, 'rf_model.pkl')
loaded = joblib.load('rf_model.pkl')
print(loaded.score(X, y))

Save PyTorch model weights

Saves only the model's learned parameters (state_dict) as a .pt file and reloads them into an identically structured model ready for inference.

Example · python
import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3))
torch.save(model.state_dict(), 'model_weights.pt')

# Reload
new_model = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3))
new_model.load_state_dict(torch.load('model_weights.pt'))
new_model.eval()
print('Loaded successfully')

Save Keras model to disk

Saves the entire Keras model architecture, weights, and compilation state to a .keras file and loads it back as a fully functional model.

Example · python
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np

model = keras.Sequential([
    layers.Dense(16, activation='relu', input_shape=(4,)),
    layers.Dense(1, activation='sigmoid')
])
model.compile(optimizer='adam', loss='binary_crossentropy')
model.fit(np.random.randn(100,4), np.random.randint(0,2,100), epochs=3, verbose=0)

model.save('my_model.keras')
loaded = keras.models.load_model('my_model.keras')
print(loaded.summary())

Discussion

  • Be the first to comment on this lesson.