Lesson 18

Saving and Loading Models

Save a trained model or pipeline with joblib or pickle, load it back to predict later, and store the details needed to use it safely afterwards.

Intermediate25 min

Saving and Loading Models hero Module 27 Lesson 18

Training can take minutes or hours, and a model is only useful if something can use it later: a web app, a scheduled job, a colleague. Saving writes the fitted object to a file. Loading gives back the same model, ready to predict without training again.

A trained pipeline saved to a file with joblib and loaded back to predict

pickle turns a model into bytes

Python's built-in pickle can serialise almost any object, fitted models included. dumps returns bytes and loads rebuilds the object.

python

import numpy as np
import pandas as pd

rng = np.random.default_rng(27)
learners = pd.DataFrame({
    "hours": rng.uniform(1, 12, 200).round(1),
    "practice": rng.integers(0, 41, 200),
    "video_minutes": rng.integers(0, 3001, 200),
    "track": rng.choice(["Python", "Data", "Web"], 200),
})
bonus = learners["track"].map({"Python": 0, "Data": 6, "Web": -5})
noise = rng.normal(0, 6, 200)
learners["score"] = (
    22 + learners["hours"] * 3.4 + learners["practice"] * 0.55 + bonus + noise
).clip(0, 100).round()
learners["passed"] = (learners["score"] >= 60).astype(int)

import pickle

from sklearn.linear_model import LogisticRegression

model = LogisticRegression().fit(learners[["hours", "practice"]], learners["passed"])
blob = pickle.dumps(model)
restored = pickle.loads(blob)

print(type(blob).__name__, type(restored).__name__)
print((restored.predict(learners[["hours", "practice"]]) == model.predict(learners[["hours", "practice"]])).all())

Output:

text

bytes LogisticRegression
True

joblib.dump and joblib.load

joblib is installed with scikit-learn and is the usual choice for models because it stores large NumPy arrays efficiently. It writes straight to a file path.

python

import numpy as np
import pandas as pd

rng = np.random.default_rng(27)
learners = pd.DataFrame({
    "hours": rng.uniform(1, 12, 200).round(1),
    "practice": rng.integers(0, 41, 200),
    "video_minutes": rng.integers(0, 3001, 200),
    "track": rng.choice(["Python", "Data", "Web"], 200),
})
bonus = learners["track"].map({"Python": 0, "Data": 6, "Web": -5})
noise = rng.normal(0, 6, 200)
learners["score"] = (
    22 + learners["hours"] * 3.4 + learners["practice"] * 0.55 + bonus + noise
).clip(0, 100).round()
learners["passed"] = (learners["score"] >= 60).astype(int)

import os

import joblib
from sklearn.ensemble import RandomForestClassifier

X = learners[["hours", "practice"]]
forest = RandomForestClassifier(random_state=0).fit(X, learners["passed"])

joblib.dump(forest, "pass_model.joblib")
print(os.path.exists("pass_model.joblib"))

loaded = joblib.load("pass_model.joblib")
print((loaded.predict(X) == forest.predict(X)).all())

Output:

text

True
True

In the playground the file lives in the browser's temporary storage and disappears when the page reloads.

Save the whole pipeline

Save the pipeline, not only the model. Then the scaler and encoder travel with it, and whoever loads it can pass raw rows.

python

import numpy as np
import pandas as pd

rng = np.random.default_rng(27)
learners = pd.DataFrame({
    "hours": rng.uniform(1, 12, 200).round(1),
    "practice": rng.integers(0, 41, 200),
    "video_minutes": rng.integers(0, 3001, 200),
    "track": rng.choice(["Python", "Data", "Web"], 200),
})
bonus = learners["track"].map({"Python": 0, "Data": 6, "Web": -5})
noise = rng.normal(0, 6, 200)
learners["score"] = (
    22 + learners["hours"] * 3.4 + learners["practice"] * 0.55 + bonus + noise
).clip(0, 100).round()
learners["passed"] = (learners["score"] >= 60).astype(int)

import joblib
from sklearn.compose import ColumnTransformer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

prepare = ColumnTransformer([
    ("numbers", StandardScaler(), ["hours", "practice", "video_minutes"]),
    ("text", OneHotEncoder(handle_unknown="ignore"), ["track"]),
])
pipe = make_pipeline(prepare, LogisticRegression())
pipe.fit(learners[["hours", "practice", "video_minutes", "track"]], learners["passed"])
joblib.dump(pipe, "pass_pipeline.joblib")

ready = joblib.load("pass_pipeline.joblib")
raw = pd.DataFrame({"hours": [8.0], "practice": [25], "video_minutes": [900], "track": ["Data"]})
print(ready.predict(raw), ready.predict_proba(raw).round(3))

Output:

text

[1] [[0.061 0.939]]

Keep the details with the model

A saved file does not explain itself. Store a small dictionary with the model: which columns it expects, what the label means, and the scikit-learn version that trained it.

python

import numpy as np
import pandas as pd

rng = np.random.default_rng(27)
learners = pd.DataFrame({
    "hours": rng.uniform(1, 12, 200).round(1),
    "practice": rng.integers(0, 41, 200),
    "video_minutes": rng.integers(0, 3001, 200),
    "track": rng.choice(["Python", "Data", "Web"], 200),
})
bonus = learners["track"].map({"Python": 0, "Data": 6, "Web": -5})
noise = rng.normal(0, 6, 200)
learners["score"] = (
    22 + learners["hours"] * 3.4 + learners["practice"] * 0.55 + bonus + noise
).clip(0, 100).round()
learners["passed"] = (learners["score"] >= 60).astype(int)

import joblib
import sklearn
from sklearn.linear_model import LogisticRegression

columns = ["hours", "practice"]
bundle = {
    "model": LogisticRegression().fit(learners[columns], learners["passed"]),
    "columns": columns,
    "label": "passed: 1 means a score of 60 or more",
    "sklearn": sklearn.__version__,
}
joblib.dump(bundle, "pass_bundle.joblib")

loaded = joblib.load("pass_bundle.joblib")
print(loaded["columns"], loaded["label"])
print("same version:", loaded["sklearn"] == sklearn.__version__)

Output:

text

['hours', 'practice'] passed: 1 means a score of 60 or more
same version: True

Load only files you trust

Loading a pickle or joblib file can run code hidden inside it, so never load one from an unknown source. Also load with the same scikit-learn version that saved it; other versions may warn or behave differently.

python

import pickle

message = {"model": "not really", "rows": 200}
blob = pickle.dumps(message)
print(pickle.loads(blob))
print("only unpickle files you created or trust")

Output:

text

{'model': 'not really', 'rows': 200}
only unpickle files you created or trust

Key takeaways

Save fitted models so they can predict later without retraining
joblib.dump and joblib.load are the usual tools for scikit-learn models
Save the whole pipeline so preprocessing travels with the model
Store the expected columns, label meaning, and version alongside the model
Only load pickle or joblib files you trust, with the same scikit-learn version
Next: What Is Machine Learning
Previous: K-Means Clustering

Loading playground…

Check your understanding

Saving and Loading Models Quiz

Why save a trained model?

1 / 10

← K-Means Clustering