Quantisation

Each weight is usually stored as a thirty two bit number, with far more precision than a network really needs. It's like measuring your height to the width of an atom. So what if we rounded every weight to one of just two hundred and fifty six values?

• Thirty two bits per weight
• Far more precision than needed
• Round to 256 values

Dynamic quantisation

This takes a single function call. We tell it which layers to convert, and that we want eight bit integers. Then we compare the size and accuracy. Run it. Four hundred kilobytes has become about a hundred. Almost four times smaller. And the accuracy is practically unchanged, and here it's even up by a tiny fraction, by chance. On a C P U, it often runs two to four times faster too. It's about the best deal in all of deep learning.

import os, torch
from base import load, net, fit, score

torch.manual_seed(0)
X, y, Xt, yt = load(4000)
model = fit(net(), X, y, epochs=3)
torch.backends.quantized.engine = "qnnpack"
small = torch.ao.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)

torch.save(model.state_dict(), "a.pt"); torch.save(small.state_dict(), "b.pt")
print(f"Size: {os.path.getsize('a.pt') / 1000:.0f} KB -> {os.path.getsize('b.pt') / 1000:.0f} KB")
print(f"Accuracy: {score(model, Xt, yt):.3f} -> {score(small, Xt, yt):.3f}")

Size: 409 KB -> 105 KB
Accuracy: 0.858 -> 0.859

import torch
from torchvision import datasets, transforms

def load(n=6000):
    t = transforms.ToTensor()
    a = datasets.MNIST("data", True, t, download=False)
    b = datasets.MNIST("data", False, t, download=False)
    return a.data[:n].float()[:, None] / 255, a.targets[:n], b.data[:2000].float()[:, None] / 255, b.targets[:2000]

def net():
    return torch.nn.Sequential(torch.nn.Flatten(), torch.nn.Linear(784, 128), torch.nn.ReLU(),
                               torch.nn.Linear(128, 10))

def fit(model, X, y, lr=1e-3, epochs=3, opt=None):
    opt = opt or torch.optim.Adam(model.parameters(), lr=lr)
    for _ in range(epochs):
        for i in torch.randperm(len(X)).split(64):
            loss = torch.nn.functional.cross_entropy(model(X[i]), y[i])
            opt.zero_grad(); loss.backward(); opt.step()
    return model

def score(model, X, y):
    return (model(X).argmax(1) == y).float().mean().item()

Choosing a method

There are three levels. Dynamic quantisation is the easiest, and it suits language models. Static quantisation also converts the activations, and needs some sample data, so it's better for vision models. Quantisation aware training simulates rounding during training, which gives the best accuracy. And for local language models, you've probably seen four bit versions, which is why they fit on a laptop.

Summary