Making a Model Cheap Enough to Ship

Week 16 of 18 · Production · 7 days

Week 16 · Production

Making a Model Cheap Enough to Ship

Week 16 · Day 1 of 7

Accuracy Is Half the Problem

Size, latency and throughput, and how to measure them honestly

Every model in this course has been judged on accuracy. That is the right measure while you are finding out whether an idea works, and it is roughly half of what decides whether the model ships. The other half is how much memory it occupies, how long one prediction takes, and what that costs multiplied by the number of requests.

This week takes one trained model and applies the four standard techniques for making it cheaper, measuring each. Two of them will do considerably less than their reputation suggests, on this model and this machine, and the week says so.

The thing being optimised

import os, time, copy
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets, transforms

tf_ = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf_)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf_)

def stack(ds, n):
    xs = torch.stack([ds[i][0] for i in range(n)])
    ys = torch.tensor([ds[i][1] for i in range(n)])
    return xs, ys

xtr, ytr = stack(train_all, 12000)
xte, yte = stack(test_all, 2000)

def accuracy(model, x=None, y=None):
    model.eval()
    x = xte if x is None else x
    y = yte if y is None else y
    with torch.no_grad():
        return (model(x).argmax(1) == y).float().mean().item()

def size_kb(model, path='tmp.pt'):
    torch.save(model.state_dict(), path)
    kb = os.path.getsize(path) / 1024
    os.remove(path)
    return kb

def latency_ms(model, n=200, batch=1):
    """Median time for one forward pass, which is what a server sees."""
    model.eval()
    x = xte[:batch]
    times = []
    with torch.no_grad():
        for _ in range(10):
            model(x)
        for _ in range(n):
            start = time.perf_counter()
            model(x)
            times.append((time.perf_counter() - start) * 1000)
    times.sort()
    return times[len(times) // 2]
def big():
    """The model we would like to ship but cannot afford to."""
    return nn.Sequential(
        nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Flatten(),
        nn.Linear(64 * 7 * 7, 256), nn.ReLU(),
        nn.Linear(256, 10))

def small():
    """Something we could actually run on a phone."""
    return nn.Sequential(
        nn.Flatten(),
        nn.Linear(784, 32), nn.ReLU(),
        nn.Linear(32, 10))

def train(model, epochs=6, lr=1e-3, seed=0, teacher=None, alpha=0.5,
          temperature=4.0):
    torch.manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    hard = nn.CrossEntropyLoss()
    soft = nn.KLDivLoss(reduction='batchmean')
    loader = DataLoader(TensorDataset(xtr, ytr), batch_size=128,
                        shuffle=True)
    for _ in range(epochs):
        model.train()
        for xb, yb in loader:
            out = model(xb)
            loss = hard(out, yb)
            if teacher is not None:
                with torch.no_grad():
                    t = teacher(xb)
                # match the teacher's whole distribution, not its answer
                kd = soft(
                    nn.functional.log_softmax(out / temperature, dim=1),
                    nn.functional.softmax(t / temperature, dim=1))
                loss = (1 - alpha) * loss + alpha * temperature ** 2 * kd
            opt.zero_grad()
            loss.backward()
            opt.step()
    model.eval()
    return model
torch.manual_seed(0)
model = train(big())
torch.save(model.state_dict(), 'teacher.pt')
print('%-22s %10s %10s %12s' % ('', 'accuracy', 'size KB', 'latency ms'))
print('%-22s %10.4f %10.1f %12.3f'
      % ('the big model', accuracy(model), size_kb(model),
         latency_ms(model)))
print('%d parameters' % sum(p.numel() for p in model.parameters()))

                         accuracy    size KB   latency ms
the big model              0.9775     3223.5        0.310
824458 parameters

Latency and throughput: Latency is how long one request waits. Throughput is how many you serve per second. They are not the same thing and improving one often costs the other, which week 17 measures directly. This week uses the median of many single-image forward passes, because that is what a user experiences.

Measure the median, and warm up first

The first few forward passes through a fresh model are slower while allocators and caches settle, so a timing loop that includes them reports a number nobody will ever see. Every measurement here runs ten passes it throws away, then takes the median of two hundred more. The mean would be dragged around by occasional scheduler interruptions.

The four techniques

Quantisation, storing weights in 8 bits instead of 32.
Pruning, setting the least important weights to zero.
Distillation, training a small model to imitate a large one.
Export, taking the model out of Python into a runtime built for inference.

They are not alternatives. The last one composes with all the others, and as it turns out it is also the one that did the most here.

Week 16 · Day 2 of 7

Quantisation

Eight bits instead of thirty-two, and the speedup that went backwards

By Pawan Rama Mali763 words

A float32 weight uses four bytes to store a number the network is largely indifferent to the fourth decimal place of. Quantisation stores it in one byte instead, with a scale factor per tensor to map the small integer range back onto the real one.

Dynamic quantisation: Weights are converted to 8 bit integers once, ahead of time. Activations are quantised on the fly as each layer runs, using a range computed from the data actually passing through. It requires no calibration data and no retraining, which makes it the one to try first.

import os, time, copy
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets, transforms

tf_ = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf_)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf_)

def stack(ds, n):
    xs = torch.stack([ds[i][0] for i in range(n)])
    ys = torch.tensor([ds[i][1] for i in range(n)])
    return xs, ys

xtr, ytr = stack(train_all, 12000)
xte, yte = stack(test_all, 2000)

def accuracy(model, x=None, y=None):
    model.eval()
    x = xte if x is None else x
    y = yte if y is None else y
    with torch.no_grad():
        return (model(x).argmax(1) == y).float().mean().item()

def size_kb(model, path='tmp.pt'):
    torch.save(model.state_dict(), path)
    kb = os.path.getsize(path) / 1024
    os.remove(path)
    return kb

def latency_ms(model, n=200, batch=1):
    """Median time for one forward pass, which is what a server sees."""
    model.eval()
    x = xte[:batch]
    times = []
    with torch.no_grad():
        for _ in range(10):
            model(x)
        for _ in range(n):
            start = time.perf_counter()
            model(x)
            times.append((time.perf_counter() - start) * 1000)
    times.sort()
    return times[len(times) // 2]
def big():
    """The model we would like to ship but cannot afford to."""
    return nn.Sequential(
        nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Flatten(),
        nn.Linear(64 * 7 * 7, 256), nn.ReLU(),
        nn.Linear(256, 10))

def small():
    """Something we could actually run on a phone."""
    return nn.Sequential(
        nn.Flatten(),
        nn.Linear(784, 32), nn.ReLU(),
        nn.Linear(32, 10))

def train(model, epochs=6, lr=1e-3, seed=0, teacher=None, alpha=0.5,
          temperature=4.0):
    torch.manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    hard = nn.CrossEntropyLoss()
    soft = nn.KLDivLoss(reduction='batchmean')
    loader = DataLoader(TensorDataset(xtr, ytr), batch_size=128,
                        shuffle=True)
    for _ in range(epochs):
        model.train()
        for xb, yb in loader:
            out = model(xb)
            loss = hard(out, yb)
            if teacher is not None:
                with torch.no_grad():
                    t = teacher(xb)
                # match the teacher's whole distribution, not its answer
                kd = soft(
                    nn.functional.log_softmax(out / temperature, dim=1),
                    nn.functional.softmax(t / temperature, dim=1))
                loss = (1 - alpha) * loss + alpha * temperature ** 2 * kd
            opt.zero_grad()
            loss.backward()
            opt.step()
    model.eval()
    return model
from torch.ao.quantization import quantize_dynamic
model = big()
model.load_state_dict(torch.load('teacher.pt'))
model.eval()
q = quantize_dynamic(copy.deepcopy(model), {nn.Linear}, dtype=torch.qint8)
print('%-22s %10s %10s %12s' % ('', 'accuracy', 'size KB', 'latency ms'))
for name, m in [('float32', model), ('int8 linear layers', q)]:
    print('%-22s %10.4f %10.1f %12.3f'
          % (name, accuracy(m), size_kb(m), latency_ms(m)))

                         accuracy    size KB   latency ms
float32                    0.9775     3223.5        0.259
int8 linear layers         0.9770      865.2        0.559

The size result is exactly as advertised: 3223 KB down to 865, a factor of 3.7, for five hundredths of a percent of accuracy. If your constraint is the size of a download or the memory of a device, that is an excellent trade and there is very little to think about.

It got slower, not faster

0.259 ms became 0.559 ms. Quantisation is usually sold as a speed improvement and here it doubled the latency.

The reason is that int8 matrix multiplication is only faster when the matrix is large enough for the arithmetic to dominate. These layers are small, so the run is dominated by converting activations to int8 on the way in and back to float on the way out, which is pure overhead. On a large language model the same change is a substantial speedup. On this one it is a regression.

That is the general shape of it. Quantisation reliably buys memory. Whether it buys speed depends on the size of your layers and whether the hardware has integer kernels worth using, and the only way to know is the measurement above.

The other kinds

	What it needs	When to use it
Dynamic	Nothing	First thing to try, especially for linear and recurrent layers
Static	A calibration pass over sample data	When you need activations quantised ahead of time too, usually for convolutional models
Quantisation aware training	A fine-tuning run	When the accuracy drop from the other two is too large

Week 16 · Day 3 of 7

Pruning

Ninety percent of the weights removed for free, and why the file did not shrink

By Pawan Rama Mali769 words

Most of the weights in a trained network are small, and a weight near zero contributes almost nothing. Pruning sets the smallest ones to exactly zero and asks what that cost.

import os, time, copy
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets, transforms

tf_ = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf_)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf_)

def stack(ds, n):
    xs = torch.stack([ds[i][0] for i in range(n)])
    ys = torch.tensor([ds[i][1] for i in range(n)])
    return xs, ys

xtr, ytr = stack(train_all, 12000)
xte, yte = stack(test_all, 2000)

def accuracy(model, x=None, y=None):
    model.eval()
    x = xte if x is None else x
    y = yte if y is None else y
    with torch.no_grad():
        return (model(x).argmax(1) == y).float().mean().item()

def size_kb(model, path='tmp.pt'):
    torch.save(model.state_dict(), path)
    kb = os.path.getsize(path) / 1024
    os.remove(path)
    return kb

def latency_ms(model, n=200, batch=1):
    """Median time for one forward pass, which is what a server sees."""
    model.eval()
    x = xte[:batch]
    times = []
    with torch.no_grad():
        for _ in range(10):
            model(x)
        for _ in range(n):
            start = time.perf_counter()
            model(x)
            times.append((time.perf_counter() - start) * 1000)
    times.sort()
    return times[len(times) // 2]
def big():
    """The model we would like to ship but cannot afford to."""
    return nn.Sequential(
        nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Flatten(),
        nn.Linear(64 * 7 * 7, 256), nn.ReLU(),
        nn.Linear(256, 10))

def small():
    """Something we could actually run on a phone."""
    return nn.Sequential(
        nn.Flatten(),
        nn.Linear(784, 32), nn.ReLU(),
        nn.Linear(32, 10))

def train(model, epochs=6, lr=1e-3, seed=0, teacher=None, alpha=0.5,
          temperature=4.0):
    torch.manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    hard = nn.CrossEntropyLoss()
    soft = nn.KLDivLoss(reduction='batchmean')
    loader = DataLoader(TensorDataset(xtr, ytr), batch_size=128,
                        shuffle=True)
    for _ in range(epochs):
        model.train()
        for xb, yb in loader:
            out = model(xb)
            loss = hard(out, yb)
            if teacher is not None:
                with torch.no_grad():
                    t = teacher(xb)
                # match the teacher's whole distribution, not its answer
                kd = soft(
                    nn.functional.log_softmax(out / temperature, dim=1),
                    nn.functional.softmax(t / temperature, dim=1))
                loss = (1 - alpha) * loss + alpha * temperature ** 2 * kd
            opt.zero_grad()
            loss.backward()
            opt.step()
    model.eval()
    return model
import torch.nn.utils.prune as prune
base = big()
base.load_state_dict(torch.load('teacher.pt'))
base.eval()
print('%10s %12s %12s' % ('sparsity', 'accuracy', 'zeros'))
for amount in [0.0, 0.5, 0.8, 0.9, 0.95, 0.99]:
    m = copy.deepcopy(base)
    targets = [(mod, 'weight') for mod in m
               if isinstance(mod, (nn.Linear, nn.Conv2d))]
    if amount > 0:
        prune.global_unstructured(targets, prune.L1Unstructured,
                                  amount=amount)
        for mod, nm in targets:
            prune.remove(mod, nm)
    total = sum(p.numel() for p in m.parameters())
    zeros = sum(int((p == 0).sum()) for p in m.parameters())
    print('%10.2f %12.4f %11.1f%%'
          % (amount, accuracy(m), 100.0 * zeros / total))

  sparsity     accuracy        zeros
00       0.9775         0.0%
50       0.9770        50.0%
80       0.9790        80.0%
90       0.9780        90.0%
95       0.9290        95.0%
99       0.4980        99.0%

Ninety percent of the weights can be deleted with no retraining at all and accuracy holds. At eighty percent it is very slightly higher than the original, which is a regularisation effect and not something to get excited about at this margin. Then it falls off a cliff: 0.929 at ninety-five percent and 0.498 at ninety-nine.

That shape is the useful part. There is a wide plateau where pruning is free, followed by a sharp edge, and the edge is not where intuition puts it. You find it by sweeping, exactly as above, and you take a setting comfortably inside the plateau rather than at its boundary.

Ninety percent zeros made the file no smaller

Look at what pruning actually did: it wrote zeros into a dense tensor. The tensor has the same shape, occupies the same memory and takes the same time to multiply. A dense zero costs exactly as much as a dense anything else.

Turning sparsity into savings needs either a sparse storage format with kernels that exploit it, or structured pruning that removes whole channels so the tensor genuinely gets smaller. Unstructured pruning on its own is a research result, not a deployment technique, and it is very often reported as though it were the second thing.

Structured pruning

Removing an entire output channel of a convolution, rather than scattered individual weights, leaves a smaller dense tensor that every existing kernel handles at full speed. It costs more accuracy per parameter removed, because you are deleting useful weights alongside useless ones, and it is the version that actually makes a model faster. If you want speed, this is the one.

Week 16 · Day 4 of 7

Distillation

Learning from a teacher's whole distribution, and a negative result checked twice

By Pawan Rama Mali1743 words

The third technique changes the model rather than compressing it. Train a small network, but instead of showing it only the correct label, show it what a large trained network believes.

import os, time, copy
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets, transforms

tf_ = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf_)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf_)

def stack(ds, n):
    xs = torch.stack([ds[i][0] for i in range(n)])
    ys = torch.tensor([ds[i][1] for i in range(n)])
    return xs, ys

xtr, ytr = stack(train_all, 12000)
xte, yte = stack(test_all, 2000)

def accuracy(model, x=None, y=None):
    model.eval()
    x = xte if x is None else x
    y = yte if y is None else y
    with torch.no_grad():
        return (model(x).argmax(1) == y).float().mean().item()

def size_kb(model, path='tmp.pt'):
    torch.save(model.state_dict(), path)
    kb = os.path.getsize(path) / 1024
    os.remove(path)
    return kb

def latency_ms(model, n=200, batch=1):
    """Median time for one forward pass, which is what a server sees."""
    model.eval()
    x = xte[:batch]
    times = []
    with torch.no_grad():
        for _ in range(10):
            model(x)
        for _ in range(n):
            start = time.perf_counter()
            model(x)
            times.append((time.perf_counter() - start) * 1000)
    times.sort()
    return times[len(times) // 2]
def big():
    """The model we would like to ship but cannot afford to."""
    return nn.Sequential(
        nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Flatten(),
        nn.Linear(64 * 7 * 7, 256), nn.ReLU(),
        nn.Linear(256, 10))

def small():
    """Something we could actually run on a phone."""
    return nn.Sequential(
        nn.Flatten(),
        nn.Linear(784, 32), nn.ReLU(),
        nn.Linear(32, 10))

def train(model, epochs=6, lr=1e-3, seed=0, teacher=None, alpha=0.5,
          temperature=4.0):
    torch.manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    hard = nn.CrossEntropyLoss()
    soft = nn.KLDivLoss(reduction='batchmean')
    loader = DataLoader(TensorDataset(xtr, ytr), batch_size=128,
                        shuffle=True)
    for _ in range(epochs):
        model.train()
        for xb, yb in loader:
            out = model(xb)
            loss = hard(out, yb)
            if teacher is not None:
                with torch.no_grad():
                    t = teacher(xb)
                # match the teacher's whole distribution, not its answer
                kd = soft(
                    nn.functional.log_softmax(out / temperature, dim=1),
                    nn.functional.softmax(t / temperature, dim=1))
                loss = (1 - alpha) * loss + alpha * temperature ** 2 * kd
            opt.zero_grad()
            loss.backward()
            opt.step()
    model.eval()
    return model
teacher = big()
teacher.load_state_dict(torch.load('teacher.pt'))
teacher.eval()
with torch.no_grad():
    probs = nn.functional.softmax(teacher(xte[:1]) / 4.0, dim=1)[0]
print('true label %d' % yte[0])
print('the teacher, softened, says:')
for digit, p in sorted(enumerate(probs.tolist()), key=lambda t: -t[1])[:5]:
    print('   %d  %.4f' % (digit, p))
print()
print('the one-hot label says 1.0 for %d and 0.0 for everything else'
      % yte[0])

true label 7
the teacher, softened, says:
   7  0.8008
   2  0.0750
   3  0.0533
   9  0.0221
   8  0.0196

the one-hot label says 1.0 for 7 and 0.0 for everything else

That is the argument in one output. The label says this is a seven and says nothing else. The teacher says it is a seven, and that if it were not, a two or a three would be the next most reasonable readings, and a zero would be absurd. The claim is that the second set carries more information per example, so a small model can learn more from the same data.

Temperature, in distillation: Dividing logits by a number above one before the softmax, which flattens the distribution and makes the small probabilities visible. Without it the teacher's output is nearly one-hot and there is nothing extra to learn from. The loss is scaled by temperature squared to keep the gradient magnitude comparable.

Measured

import os, time, copy
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets, transforms

tf_ = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf_)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf_)

def stack(ds, n):
    xs = torch.stack([ds[i][0] for i in range(n)])
    ys = torch.tensor([ds[i][1] for i in range(n)])
    return xs, ys

xtr, ytr = stack(train_all, 12000)
xte, yte = stack(test_all, 2000)

def accuracy(model, x=None, y=None):
    model.eval()
    x = xte if x is None else x
    y = yte if y is None else y
    with torch.no_grad():
        return (model(x).argmax(1) == y).float().mean().item()

def size_kb(model, path='tmp.pt'):
    torch.save(model.state_dict(), path)
    kb = os.path.getsize(path) / 1024
    os.remove(path)
    return kb

def latency_ms(model, n=200, batch=1):
    """Median time for one forward pass, which is what a server sees."""
    model.eval()
    x = xte[:batch]
    times = []
    with torch.no_grad():
        for _ in range(10):
            model(x)
        for _ in range(n):
            start = time.perf_counter()
            model(x)
            times.append((time.perf_counter() - start) * 1000)
    times.sort()
    return times[len(times) // 2]
def big():
    """The model we would like to ship but cannot afford to."""
    return nn.Sequential(
        nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Flatten(),
        nn.Linear(64 * 7 * 7, 256), nn.ReLU(),
        nn.Linear(256, 10))

def small():
    """Something we could actually run on a phone."""
    return nn.Sequential(
        nn.Flatten(),
        nn.Linear(784, 32), nn.ReLU(),
        nn.Linear(32, 10))

def train(model, epochs=6, lr=1e-3, seed=0, teacher=None, alpha=0.5,
          temperature=4.0):
    torch.manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    hard = nn.CrossEntropyLoss()
    soft = nn.KLDivLoss(reduction='batchmean')
    loader = DataLoader(TensorDataset(xtr, ytr), batch_size=128,
                        shuffle=True)
    for _ in range(epochs):
        model.train()
        for xb, yb in loader:
            out = model(xb)
            loss = hard(out, yb)
            if teacher is not None:
                with torch.no_grad():
                    t = teacher(xb)
                # match the teacher's whole distribution, not its answer
                kd = soft(
                    nn.functional.log_softmax(out / temperature, dim=1),
                    nn.functional.softmax(t / temperature, dim=1))
                loss = (1 - alpha) * loss + alpha * temperature ** 2 * kd
            opt.zero_grad()
            loss.backward()
            opt.step()
    model.eval()
    return model
teacher = big()
teacher.load_state_dict(torch.load('teacher.pt'))
teacher.eval()
print('%-30s %10s %10s %12s'
      % ('', 'accuracy', 'size KB', 'latency ms'))
print('%-30s %10.4f %10.1f %12.3f'
      % ('teacher', accuracy(teacher), size_kb(teacher),
         latency_ms(teacher)))
alone = train(small(), seed=1)
print('%-30s %10.4f %10.1f %12.3f'
      % ('student, labels only', accuracy(alone), size_kb(alone),
         latency_ms(alone)))
taught = train(small(), seed=1, teacher=teacher)
print('%-30s %10.4f %10.1f %12.3f'
      % ('student, distilled', accuracy(taught), size_kb(taught),
         latency_ms(taught)))

                                 accuracy    size KB   latency ms
teacher                            0.9775     3223.5        0.271
student, labels only               0.8905      101.6        0.028
student, distilled                 0.8730      101.6        0.031

The student is 32 times smaller and roughly 9 times faster, for about nine points of accuracy. That part is the real trade and it is often worth taking. But distillation made the student worse than training it on labels alone, which is the opposite of the entire point.

Checking that before believing it

One run at one setting is not enough to publish a negative result on, so here is a sweep over the mixing weight, at two seeds each.

import os, time, copy
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets, transforms

tf_ = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf_)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf_)

def stack(ds, n):
    xs = torch.stack([ds[i][0] for i in range(n)])
    ys = torch.tensor([ds[i][1] for i in range(n)])
    return xs, ys

xtr, ytr = stack(train_all, 12000)
xte, yte = stack(test_all, 2000)

def accuracy(model, x=None, y=None):
    model.eval()
    x = xte if x is None else x
    y = yte if y is None else y
    with torch.no_grad():
        return (model(x).argmax(1) == y).float().mean().item()

def size_kb(model, path='tmp.pt'):
    torch.save(model.state_dict(), path)
    kb = os.path.getsize(path) / 1024
    os.remove(path)
    return kb

def latency_ms(model, n=200, batch=1):
    """Median time for one forward pass, which is what a server sees."""
    model.eval()
    x = xte[:batch]
    times = []
    with torch.no_grad():
        for _ in range(10):
            model(x)
        for _ in range(n):
            start = time.perf_counter()
            model(x)
            times.append((time.perf_counter() - start) * 1000)
    times.sort()
    return times[len(times) // 2]
def big():
    """The model we would like to ship but cannot afford to."""
    return nn.Sequential(
        nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Flatten(),
        nn.Linear(64 * 7 * 7, 256), nn.ReLU(),
        nn.Linear(256, 10))

def small():
    """Something we could actually run on a phone."""
    return nn.Sequential(
        nn.Flatten(),
        nn.Linear(784, 32), nn.ReLU(),
        nn.Linear(32, 10))

def train(model, epochs=6, lr=1e-3, seed=0, teacher=None, alpha=0.5,
          temperature=4.0):
    torch.manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    hard = nn.CrossEntropyLoss()
    soft = nn.KLDivLoss(reduction='batchmean')
    loader = DataLoader(TensorDataset(xtr, ytr), batch_size=128,
                        shuffle=True)
    for _ in range(epochs):
        model.train()
        for xb, yb in loader:
            out = model(xb)
            loss = hard(out, yb)
            if teacher is not None:
                with torch.no_grad():
                    t = teacher(xb)
                # match the teacher's whole distribution, not its answer
                kd = soft(
                    nn.functional.log_softmax(out / temperature, dim=1),
                    nn.functional.softmax(t / temperature, dim=1))
                loss = (1 - alpha) * loss + alpha * temperature ** 2 * kd
            opt.zero_grad()
            loss.backward()
            opt.step()
    model.eval()
    return model
teacher = big()
teacher.load_state_dict(torch.load('teacher.pt'))
teacher.eval()
print('alpha 0 is labels only, alpha 1 is the teacher only')
print('%8s %12s %12s' % ('alpha', 'seed 1', 'seed 2'))
for alpha in [0.0, 0.3, 0.5, 0.7, 0.9]:
    accs = []
    for seed in [1, 2]:
        t = None if alpha == 0 else teacher
        m = train(small(), seed=seed, teacher=t, alpha=alpha)
        accs.append(accuracy(m))
    print('%8.1f %12.4f %12.4f' % (alpha, accs[0], accs[1]))

alpha 0 is labels only, alpha 1 is the teacher only
   alpha       seed 1       seed 2
0       0.8905       0.8920
3       0.8835       0.8795
5       0.8835       0.8800
7       0.8835       0.8785
9       0.8820       0.8775

Consistent, monotonic, and in the wrong direction

Every alpha above zero is worse than alpha zero, at both seeds, and it gets worse the more weight the teacher is given. This is not noise and it is not one unlucky hyperparameter.

The most likely explanation is capacity. A 784 to 32 to 10 network has nowhere near enough room to represent the teacher's function, so asking it to match a full ten-way distribution spends its very limited capacity on relationships it cannot use, at the expense of the decision boundary it was actually going to be scored on. Distillation helps when the student is small but not hopeless.

The lesson is not that distillation does not work. It works well enough that most deployed small language models are made this way. The lesson is that it is a technique with a range of validity, and that the honest way to find out whether you are inside it is two training runs and a table.

Week 16 · Day 5 of 7

Getting Out of Python

ONNX export, the parity check, and the largest win of the week

By Pawan Rama Mali783 words

The fourth technique does not change the model at all. It takes the model out of Python.

A PyTorch model in eager mode dispatches every operation through the interpreter, one at a time, checking types and shapes as it goes. For training that flexibility is the whole point. For inference, where the graph never changes, it is pure overhead.

ONNX: A framework-neutral description of a computation graph. Both PyTorch and Keras can export to it, and runtimes exist for servers, phones, browsers and microcontrollers. Exporting is how a model stops being a Python object and becomes an artefact something else can run.

import os, time, copy
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets, transforms

tf_ = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf_)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf_)

def stack(ds, n):
    xs = torch.stack([ds[i][0] for i in range(n)])
    ys = torch.tensor([ds[i][1] for i in range(n)])
    return xs, ys

xtr, ytr = stack(train_all, 12000)
xte, yte = stack(test_all, 2000)

def accuracy(model, x=None, y=None):
    model.eval()
    x = xte if x is None else x
    y = yte if y is None else y
    with torch.no_grad():
        return (model(x).argmax(1) == y).float().mean().item()

def size_kb(model, path='tmp.pt'):
    torch.save(model.state_dict(), path)
    kb = os.path.getsize(path) / 1024
    os.remove(path)
    return kb

def latency_ms(model, n=200, batch=1):
    """Median time for one forward pass, which is what a server sees."""
    model.eval()
    x = xte[:batch]
    times = []
    with torch.no_grad():
        for _ in range(10):
            model(x)
        for _ in range(n):
            start = time.perf_counter()
            model(x)
            times.append((time.perf_counter() - start) * 1000)
    times.sort()
    return times[len(times) // 2]
def big():
    """The model we would like to ship but cannot afford to."""
    return nn.Sequential(
        nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Flatten(),
        nn.Linear(64 * 7 * 7, 256), nn.ReLU(),
        nn.Linear(256, 10))

def small():
    """Something we could actually run on a phone."""
    return nn.Sequential(
        nn.Flatten(),
        nn.Linear(784, 32), nn.ReLU(),
        nn.Linear(32, 10))

def train(model, epochs=6, lr=1e-3, seed=0, teacher=None, alpha=0.5,
          temperature=4.0):
    torch.manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    hard = nn.CrossEntropyLoss()
    soft = nn.KLDivLoss(reduction='batchmean')
    loader = DataLoader(TensorDataset(xtr, ytr), batch_size=128,
                        shuffle=True)
    for _ in range(epochs):
        model.train()
        for xb, yb in loader:
            out = model(xb)
            loss = hard(out, yb)
            if teacher is not None:
                with torch.no_grad():
                    t = teacher(xb)
                # match the teacher's whole distribution, not its answer
                kd = soft(
                    nn.functional.log_softmax(out / temperature, dim=1),
                    nn.functional.softmax(t / temperature, dim=1))
                loss = (1 - alpha) * loss + alpha * temperature ** 2 * kd
            opt.zero_grad()
            loss.backward()
            opt.step()
    model.eval()
    return model
import numpy as np, onnxruntime as ort
model = big()
model.load_state_dict(torch.load('teacher.pt'))
model.eval()
# dynamo=False selects the long-standing tracing exporter, which needs
# no extra packages. The newer one requires onnxscript.
torch.onnx.export(model, (xte[:1],), 'm.onnx', dynamo=False,
                  input_names=['image'], output_names=['logits'],
                  dynamic_axes={'image': {0: 'batch'},
                                'logits': {0: 'batch'}})
sess = ort.InferenceSession('m.onnx',
                            providers=['CPUExecutionProvider'])
onnx_out = sess.run(None, {'image': xte[:256].numpy()})[0]
with torch.no_grad():
    torch_out = model(xte[:256]).numpy()
print('largest disagreement %.2e'
      % float(np.abs(onnx_out - torch_out).max()))
print('same predictions: %s'
      % bool((onnx_out.argmax(1) == torch_out.argmax(1)).all()))
print('file size %.1f KB' % (os.path.getsize('m.onnx') / 1024))
print()
one = xte[:1].numpy()
for _ in range(10):
    sess.run(None, {'image': one})
times = []
for _ in range(200):
    start = time.perf_counter()
    sess.run(None, {'image': one})
    times.append((time.perf_counter() - start) * 1000)
times.sort()
print('%-22s %10.3f ms' % ('onnxruntime', times[len(times) // 2]))
print('%-22s %10.3f ms' % ('pytorch eager', latency_ms(model)))

largest disagreement 3.81e-06
same predictions: True
file size 3221.9 KB

onnxruntime                 0.117 ms
pytorch eager               0.509 ms

Four times faster, with predictions that are identical and numbers that agree to about four parts in a million. No accuracy was traded, no retraining happened, and the model was not modified. Of the four techniques in this week, this is the one that did the most and asked for the least.

Always check the exported model against the original

Export works by tracing: the exporter runs your model once and records the operations. Anything that depended on the specific input is baked in. A branch on a tensor value takes whichever path that example took, a shape computed from the data becomes a constant, and the exported graph then silently disagrees with your model on every other input.

The disagreement check above is three lines and it is the difference between finding that now and finding it from production traffic. Run it on a batch, not one example, and use dynamic_axes so the batch dimension is not frozen at whatever you traced with.

Week 16 · Day 6 of 7

All Four, Side by Side

The comparison table, and the order to try them in

By Pawan Rama Mali410 words

Four techniques, one model, all measured on the same machine. Here is the whole picture.

Technique	Accuracy	Size	Latency	Verdict
None	0.9775	3223 KB	0.31 ms	The baseline
Dynamic quantisation	0.9770	865 KB	0.56 ms	3.7x smaller, and slower
Pruning to 90 percent	0.9780	3223 KB	unchanged	No saving at all without sparse kernels
Distillation to a small model	0.8730	102 KB	0.03 ms	32x smaller, 9 points worse, and worse than the same student trained on labels
ONNX export	0.9775	3222 KB	0.12 ms	4x faster, nothing given up

Read that table as a warning about where advice comes from. Three of these four techniques are described everywhere as ways to make models faster. On this model, on this hardware, one of them was slower, one changed nothing measurable, and the one that worked is the one that is least often described as an optimisation at all.

Which is not to say the others are useless

Quantisation is the right answer when memory is the constraint, and it genuinely does accelerate large matrix multiplications, which is why every deployed language model uses it.
Pruning becomes real when it is structured, or when the runtime has sparse kernels. The plateau it revealed, ninety percent of weights removable for nothing, is a true fact about the model regardless.
Distillation is how most small production models are made. It needs a student with enough capacity to benefit, which this one did not.

The order to try them in

Export first. It is the cheapest to try, it composes with everything else, and it may end the exercise.
Then ask which resource is actually binding. Memory, latency and throughput want different techniques and people routinely optimise the wrong one.
Quantise if the answer is memory.
Reach for a smaller architecture, distilled if it helps, when export and quantisation are not enough. This is the only route to a large speedup and the only one that costs real accuracy.
Measure every step against the unmodified model on the same machine, and keep the accuracy check in the loop.

The rule this week is really about

Every one of these techniques has a regime where it helps and a regime where it does nothing or hurts, and which regime you are in depends on your model and your hardware rather than on the technique. The measurements take minutes. Taking the advice without them is how a team ships something slower than what they started with.

Week 16 · Day 7 of 7

The Decisions That Matter More

Architecture, resolution, and profiling the whole request

By Pawan Rama Mali417 words

Everything above optimised a model that had already been chosen. The larger decisions happen earlier, and they matter more.

The decisions that dominate

Model size. Nothing in this week recovered the difference between the big model and the small one. Choosing an architecture that fits the budget beats compressing one that does not.
Input size. Convolutional cost scales with the number of pixels. Halving the resolution is close to a four times saving and is frequently free in accuracy, and it is tried far less often than quantisation.
Whether the model runs at all. Caching repeated requests, or a cheap filter that handles the easy cases and escalates only the hard ones, changes the arithmetic more than any of this.

Where the cost really lives in production

The measurements in this week are all of the forward pass. In a real service the forward pass is often not the largest term. Decoding a JPEG, resizing it, moving bytes over the network and serialising a response can each take longer than the model, and week 17 measures a service end to end rather than a model in isolation.

Profile before optimising, and profile the whole thing

A team that spends a week quantising a model whose latency is eighty percent image decoding has made the product four percent faster. This is the single most common way effort is wasted in deployment, and it is entirely avoidable by measuring the whole request path once before choosing what to work on.

The checklist

Measure the baseline: accuracy, size, and median latency with a warmup.
Export to ONNX and check the outputs match on a batch. Keep the speedup if the check passes.
Decide which resource is binding before choosing a technique.
Quantise for memory. Verify accuracy and re-measure latency, because it can go the wrong way.
Sweep pruning to find the plateau edge, and be honest about whether your runtime can exploit sparsity at all.
Consider a smaller architecture, and try distillation with at least two seeds before concluding anything about it.
Re-measure everything together at the end. The techniques interact, and the combination is not the sum of the parts.

What week 17 does with this

A model that is small and fast is still not a product. The next week puts one behind an interface, measures what a request actually costs, breaks the preprocessing on purpose to see what that does, and works out whether you could tell your model had started failing if nobody sent you labels.