Making Models Smaller

A model that is too big, or too slow, costs money and annoys users. Sometimes you need it to run on a phone, or on a small server. The good news is that models are often much bigger than they need to be. In this lesson, you will shrink one, and check how much quality you lose, which is often, almost none.

• Big models cost money
• Shrink without losing much
• Always check the quality

Four ways to shrink a model

There are four main techniques. Quantization stores the numbers with fewer bits, for example, eight bit integers, instead of thirty two bit decimals. Pruning removes weights, or whole neurons, that contribute little. Distillation trains a small student model, to imitate a large teacher, so you keep most of the quality, at a fraction of the size. And choosing a smaller architecture in the first place. Quantization is the easiest to try, because it needs no retraining.

Quantizing a model

Here is quantization, on a real network. We train a neural network, with two hidden layers, and export it to O N N X, in the normal thirty two bit form. Then one call quantizes the weights to eight bit integers. We load both files, run the same test customers through each, and print the file size, and the accuracy. The compressed model is about a third of the size, and its accuracy is essentially identical. Smaller files load faster, use less memory, and often run faster, too.

import os
import numpy as np
import onnxruntime as ort
from onnxruntime.quantization import quantize_dynamic, QuantType
from skl2onnx import to_onnx
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import accuracy_score
from churn import make_data

df = make_data(4000, seed=1)
X, y = df.drop(columns="churn").astype(np.float32), df["churn"]
model = MLPClassifier(hidden_layer_sizes=(128, 128), max_iter=300, random_state=0).fit(X[:3000], y[:3000])

with open("model_fp32.onnx", "wb") as f:
    f.write(to_onnx(model, X.to_numpy()[:1], options={"zipmap": False}).SerializeToString())
quantize_dynamic("model_fp32.onnx", "model_int8.onnx", weight_type=QuantType.QInt8)

for name in ("model_fp32.onnx", "model_int8.onnx"):
    session = ort.InferenceSession(name)
    labels = session.run(None, {session.get_inputs()[0].name: X.to_numpy()[3000:]})[0]
    size_kb = os.path.getsize(name) / 1024
    print(f"{name:16s} {size_kb:7.1f} KB   accuracy {accuracy_score(y[3000:], labels):.3f}")

model_fp32.onnx     68.9 KB   accuracy 0.834
model_int8.onnx     21.5 KB   accuracy 0.836

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})

The trade-offs of compression

Compression is not free, so be careful. Quantization may reduce accuracy a little, and sometimes a lot, for sensitive models. Pruning and distillation need extra training work. The savings depend on your hardware, because some chips run eight bit numbers much faster, and others do not. So follow a simple routine. Measure the baseline. Apply one technique. Measure quality, size, and speed again, on your own data. Keep the change, only if the trade is worth it.

Summary