Packaging Models for Serving

A checkpoint is not a servable model. Serving needs the weights, the fitted preprocessing, the label map, the input schema, and the runtime versions — packaged so the server loads them as one unit. The classic production failure is a model that scores well offline because it receives different inputs online.

The idea

Package these together, always:

• Weights — a state dict, SavedModel, or ONNX graph.
• Fitted preprocessing — tokenizer, scaler mean/std, vocabulary, image normalization constants.
• Output semantics — the class index to name mapping and what the score means.
• Input schema — dtypes, shapes, which fields are optional, and null handling.
• Runtime — framework version, CPU or GPU, and the forward code itself.

Format choices trade convenience against portability:

Whatever the format, declare the signature so the server validates input before inference, and store version metadata next to the weights so a running process can report what it is serving.

Worked example

A scaler is fitted on training data and yields mean = 12.4, std = 3.1. The team saves only the weights and recomputes the scaler at serving time from the incoming request batch.

Now the same customer request returns a different score depending on what else arrived in the same batch, because the batch mean changed. Aggregate accuracy barely moves, so nothing alerts — but individual decisions are unstable and untraceable. The fix is to store the fitted scaler inside the artifact and never refit at serving.

The cheapest guard is a golden input: one fixed feature row whose expected output is recorded from the training run. A test asserts the served score matches within a tolerance, which catches normalization and label-order regressions immediately.

In code

import torch

torch.save(
    {
        "state_dict": model.state_dict(),
        "input_schema": {"x": {"dtype": "float32", "shape": [None, 32]}},
        "preprocess": {"mean": prep["mean"], "std": prep["std"]},
        "labels": LABELS,
        "framework": f"torch=={torch.__version__}",
    },
    "artifacts/model.pt",
)

golden = torch.tensor([GOLDEN_ROW])  # tests/test_serving.py, a recorded fixture          # fixture recorded from the training run
assert abs(serve(golden) - GOLDEN_EXPECTED) < 1e-6

Check yourself

1. Name three things a bare checkpoint lacks that serving requires.
2. Why is refitting preprocessing at request time a correctness bug, not only a performance issue?
3. When would you choose ONNX over a framework pickle?

Key takeaways

• Package weights with preprocessing, schema, and labels as a single unit.
• Fitted transforms are part of the model; never recompute them at serving time.
• A golden-input test is the cheapest defense against train/serve skew.