Saving & Loading Models

Understanding Saving & Loading Models

Saving a model means preserving the information needed to reproduce its predictions later. Training computes learned state from examples; inference applies that state to new inputs. For linear regression, learned state includes coefficients and an intercept. For a decision tree, it includes split rules and leaf values. Loading restores that information; it should not silently train a new model or change the prediction rule.

Preprocessing is part of that rule. Suppose training first transforms a feature using z = (x − mean) / scale, then predicts y = weight × z + intercept. The mean and scale are learned state too. Saving only the weight and intercept loses the mapping from raw inputs to predictions. In scikit-learn, a fitted Pipeline bundles transformations and the estimator so prediction applies them in the correct order.

Split data before learning preprocessing parameters. Fit the pipeline on training data only, then use its predict method on held-out data. During prediction, transformations reuse their fitted statistics rather than recomputing them. Including test rows when estimating a mean, selecting features, or learning categories leaks information into training. Persistence preserves whatever you fitted, including mistakes; saving a pipeline does not retroactively prevent leakage.

The persistence format determines what loading requires. Joblib conveniently stores scikit-learn objects, but uses pickle-based mechanisms that can execute code during loading. Never load an untrusted pickle or joblib artifact. A familiar filename, extension, or successful antivirus scan is not a trust guarantee. Keras provides the .keras format; PyTorch commonly saves a state_dict, which requires compatible model architecture code. Preprocessing outside those models still needs explicit preservation.

A saved artifact is not a complete environment. Record dependency versions, feature names and order, units, target meaning, and the code revision. Scikit-learn does not support loading persisted models across different scikit-learn versions. Recreate the training environment rather than assuming compatibility. The example below uses JSON for a tiny, explicitly defined prediction rule, avoiding executable object deserialization. This improves inspectability but requires maintaining reconstruction code and validation yourself. It is not a general replacement for pipeline persistence, nor does successful loading establish model accuracy or suitability for changed data.

Practical walkthrough

Start by defining the prediction contract: this example accepts one finite numeric feature called x and returns a numeric prediction. Training rows are deliberately fixed, and inference inputs are separate. There is no random split because this exercise tests persistence, not predictive quality. Python's standard library is sufficient; no packages or downloads are required.

Compute the training mean and population standard deviation, then standardize only the training inputs. Fit ordinary least squares in that transformed coordinate system. With one feature, the slope is the sum of centered feature–target products divided by the sum of squared centered features. The intercept adjusts for the transformed feature mean. Constant training features would make this calculation undefined, so the example assumes nonzero variance.

Build one artifact containing both preprocessing statistics and regression parameters. Add a format version and feature name so loading code can reject an incompatible structure. For a production artifact, also record units, training-data identity, dependency versions where applicable, and evaluation context. Do not include sensitive training records merely for convenience.

Write the artifact to a temporary JSON file, read it back, and check its supported format and parameter constraints. JSON parsing does not execute embedded Python objects, although external input still requires size limits and rigorous schema validation. The checks here cover this small, locally generated artifact.

Finally, predict with the original and restored state on identical inputs. Compare unrounded results with a numerical tolerance, and round only printed output. Keep this round-trip check distinct from held-out evaluation: matching predictions verify restoration, not generalization.

Worked example

import json
import math
import tempfile
from pathlib import Path

# Only these rows determine preprocessing and model parameters.
x_train = [0.0, 2.0, 4.0]
y_train = [1.0, 5.0, 9.0]
x_new = [1.0, 3.0]

mean = sum(x_train) / len(x_train)
scale = math.sqrt(sum((x - mean) ** 2 for x in x_train) / len(x_train))
z = [(x - mean) / scale for x in x_train]
z_mean = sum(z) / len(z)
y_mean = sum(y_train) / len(y_train)
weight = sum((a - z_mean) * (b - y_mean) for a, b in zip(z, y_train)) / sum((a - z_mean) ** 2 for a in z)
intercept = y_mean - weight * z_mean

artifact = {
    'format_version': 1,
    'feature': 'x',
    'mean': mean,
    'scale': scale,
    'weight': weight,
    'intercept': intercept,
}

def predict(state, values):
    return [state['weight'] * ((x - state['mean']) / state['scale'])
            + state['intercept'] for x in values]

before = predict(artifact, x_new)
with tempfile.TemporaryDirectory() as directory:
    path = Path(directory) / 'model.json'
    path.write_text(json.dumps(artifact, allow_nan=False), encoding='utf-8')
    loaded = json.loads(path.read_text(encoding='utf-8'))
    if loaded['format_version'] != 1 or loaded['feature'] != 'x':
        raise ValueError('Unsupported model contract')
    for key in ('mean', 'scale', 'weight', 'intercept'):
        if type(loaded[key]) not in (int, float) or not math.isfinite(loaded[key]):
            raise ValueError('Invalid numeric parameter')
    if loaded['scale'] <= 0:
        raise ValueError('Scale must be positive')
    after = predict(loaded, x_new)

print('Before:', [round(value, 6) for value in before])
print('After:', [round(value, 6) for value in after])
print('Match:', all(math.isclose(a, b, rel_tol=1e-12, abs_tol=1e-12)
                    for a, b in zip(before, after)))

Expected output

Before: [3.0, 7.0]
After: [3.0, 7.0]
Match: True

The training points follow y = 2x + 1 exactly. Standardization changes the coordinate system, so the fitted weight is not simply 2: it incorporates the training scale. The stored mean is 2, and the fitted intercept in standardized coordinates is 5. Restoring all four numeric parameters reconstructs the same raw-input prediction function. Inputs 1 and 3 therefore produce 3 and 7 before and after loading. The file stores data rather than a callable; predict supplies the implementation. This deliberately narrow design illustrates pipeline state without pretending that arbitrary scikit-learn estimators can be reconstructed from these few fields.

Common mistakes and how to avoid them

• Saving only the final estimator.

  Persist fitted preprocessing alongside it, preferably as one pipeline. Include category mappings, imputation values, and feature-selection state wherever those operations affect predictions.

• Calling fit or fit_transform on incoming inference batches.

  Reuse the restored transformations through predict or transform. Refitting changes the coordinate system and can make predictions depend on other rows in the batch.

• Loading a downloaded joblib file to inspect whether it is safe.

  Establish provenance before loading. Loading itself can execute code; inspecting the resulting object is too late. Use only trusted artifacts in a controlled environment.

• Treating a successful load as sufficient verification.

  Check the input contract, environment, and reference predictions. Then evaluate predictive quality separately on held-out data representative of intended use.

Review questions

Further reading

• scikit-learn user guide