Model Formats

How do you save a model to a file? The obvious answer, pickle, is easy, but has hidden problems. It ties the model to one language, and to exact library versions, and loading an untrusted pickle can run harmful code. In this lesson, you will see the options, and a portable standard called O N N X.

• Saving a model is not trivial
• Pickle has hidden dangers
• ONNX is portable

Ways to save a model

Here are the common options. Pickle and joblib are easy, Python only, and fragile across library versions, and unsafe to load from strangers. Framework formats, like the ones of Tensor Flow and Py Torch, are better for deep learning, but tied to their framework. O N N X, the open neural network exchange, is a neutral format, that many languages and runtimes can load, with a fast engine to run it. And M L flow wraps any of them in a standard package, with its own metadata.

Exporting to ONNX

Here is a model exported to O N N X. We train a simple model as usual. Then one call converts it to the O N N X format, and we save the file. Now a completely different engine, O N N X Runtime, loads that file. We feed it five customers, and print its answers. Next to them, the original answers from scikit learn. They match exactly.

import numpy as np
import onnxruntime as ort
from skl2onnx import to_onnx
from sklearn.linear_model import LogisticRegression
from churn import make_data

df = make_data(2000, seed=1)
X, y = df.drop(columns="churn").astype(np.float32), df["churn"]
model = LogisticRegression(max_iter=500).fit(X, y)

onnx_model = to_onnx(model, X.to_numpy()[:1], options={"zipmap": False})
with open("churn.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

session = ort.InferenceSession("churn.onnx")
name = session.get_inputs()[0].name
labels, probs = session.run(None, {name: X.to_numpy()[:5]})
print("onnx    :", labels.tolist())
print("sklearn :", model.predict(X.head(5)).tolist())
print("same answers:", bool((labels == model.predict(X.head(5))).all()))

onnx    : [0, 0, 0, 0, 1]
sklearn : [0, 0, 0, 0, 1]
same answers: True

import numpy as np, pandas as pd

def make_data(n, seed, shifted=False):
    rng = np.random.default_rng(seed)
    tenure = rng.integers(1, 72, n)
    charges = rng.normal(65 + (28 if shifted else 0), 18, n).clip(20, 140).round(2)
    calls = rng.poisson(1.5 + (1.5 if shifted else 0), n)
    contract = rng.choice([0, 1, 2], n, p=[.55, .25, .20] if not shifted else [.30, .30, .40])
    logit = -1.0 - 0.04 * tenure + 0.025 * (charges - 65) + 0.45 * calls - 0.8 * contract
    if shifted:
        logit = -0.2 - 0.01 * tenure - 0.01 * (charges - 65) + 0.2 * calls - 0.2 * contract
    churn = (rng.random(n) < 1 / (1 + np.exp(-logit))).astype(int)
    return pd.DataFrame({"tenure": tenure, "charges": charges, "calls": calls, "contract": contract, "churn": churn})

When to use ONNX

O N N X is valuable in a few situations. When the serving system is not written in Python, for example Java, or C sharp, or in a browser. When you need faster inference, because O N N X Runtime optimises the model graph. And when you want one stable file format, across training library versions. The cost is an extra conversion step, and the need to test that the converted model agrees with the original, exactly as we just did.

Summary