Learn AI Series):You've trained a beautiful model. It scores 95% accuracy, the validation curve is a thing of art, and you are (rightly) proud. Then you deploy it, a real user hits the endpoint, and it takes 2 full seconds to answer. Your product owner asked for 50 milliseconds. Wowzers.
That distance -- between a model that's correct and a model that's usable -- is what I call the optimization gap, and closing it is a craft of its own. At the end of episode #119 I promised we'd stop worrying about whether a model is right and start worrying about whether it's fast enough to actually reach a user. Here we are. We poked at quantization once before, back in #70 when we were squeezing local LLMs onto ordinary hardware. Today we open the whole toolbox: quantization, pruning, distillation, portable runtimes, and -- crucially -- how to measure any of it without lying to yourself.
Having said that, one rule floats above everything below: fast is only useful if it's still correct. Every technique here trades a sliver of accuracy for a slab of speed. The whole game is knowing how big a sliver you can afford, and measuring it honestly instead of hoping.
Neural network weights live as 32-bit floating-point numbers by default. That's 4 bytes per parameter, so a 100M-parameter model burns 400MB just to store its weights, never mind running them. Quantization shrinks that by representing those numbers with fewer bits -- typically INT8, sometimes INT4. The insight that makes it work: those 32 bits were always wild overkill. A weight of 0.037182946 carries far more precision than the network actually needs to make the same decision.
import torch
import torch.nn as nn
import time
class SimpleClassifier(nn.Module):
def __init__(self, input_dim=784, hidden=512, n_classes=10):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(input_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, n_classes),
)
def forward(self, x):
return self.layers(x)
# Train your model normally in FP32...
model = SimpleClassifier()
# Post-training dynamic quantization (the simplest thing that works)
quantized_model = torch.quantization.quantize_dynamic(
model,
{nn.Linear}, # which layer types to quantize
dtype=torch.qint8, # target precision
)
def model_size_mb(model):
param_bytes = sum(p.nelement() * p.element_size() for p in model.parameters())
buffer_bytes = sum(b.nelement() * b.element_size() for b in model.buffers())
return (param_bytes + buffer_bytes) / 1024 / 1024
print(f"FP32 model: {model_size_mb(model):.2f} MB")
print(f"INT8 model: {model_size_mb(quantized_model):.2f} MB")
Dynamic quantization converts the weights to INT8 up front, then quantizes the activations on the fly during inference based on the range they actually take. That "dynamic" part is why it needs zero calibration data -- it figures out the activation range at runtime. You typically get a 2-4x size reduction and a 1.5-3x speedup on CPU, at a cost of usually less than 1% accuracy. For a Linear-and-LSTM-heavy model it's genuinely a one-liner win.
When you want more, reach for static quantization, which pins down the activation ranges ahead of time using a bit of representative data:
# Static quantization: calibrate with representative data
model.eval()
# 1. Attach a quantization config describing HOW to quantize
model.qconfig = torch.quantization.get_default_qconfig('x86')
prepared = torch.quantization.prepare(model)
# 2. Calibrate: push representative batches through the model.
# This observes the real activation ranges (scale, zero_point).
with torch.no_grad():
for batch in calibration_loader:
prepared(batch)
# 3. Convert to a genuinely quantized model
quantized_static = torch.quantization.convert(prepared)
Static quantization is more accurate and often faster than dynamic, because it bakes the activation ranges into the model instead of recomputing them every forward pass. The price is that you need a few hundred representative samples to calibrate on. Nota bene: "representative" is doing real work in that sentence -- calibrate on data that looks nothing like production traffic and your scales will be wrong, quietly, in a way no exception ever warns you about. This is the same lesson as always in this series: garbage calibration in, garbage model out, and the model won't have the decency to crash about it. If your production traffic is, say, mostly night-time images and you calibrated on daylight, the activation ranges you measured are a polite fiction, and your INT8 model will drift in exactly the cases you care about most.
Sometimes post-training quantization draws too much blood -- you convert to INT8 and accuracy falls off a cliff. That tends to happen with small models (they have no spare capacity to lose) or with activation distributions full of nasty outliers. The fix is quantization-aware training: you simulate the rounding error during training, so the network learns weights that are robust to being squashed into INT8 later.
# Quantization-Aware Training
model = SimpleClassifier()
model.train()
# Insert "fake quantization" observers that round in the forward pass
model.qconfig = torch.quantization.get_default_qat_qconfig('x86')
prepared = torch.quantization.prepare_qat(model)
# Fine-tune with fake quantization active (simulates INT8 error)
optimizer = torch.optim.Adam(prepared.parameters(), lr=1e-4)
for epoch in range(5): # a handful of epochs is usually enough
for batch_x, batch_y in train_loader:
output = prepared(batch_x)
loss = nn.functional.cross_entropy(output, batch_y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Convert the fine-tuned model to real INT8
prepared.eval()
quantized_qat = torch.quantization.convert(prepared)
The clever trick hiding inside "fake quantization" is that the forward pass rounds the values (so the loss feels the precision loss), but the backward pass pretends the rounding was the identity function -- the so-called straight-through estimator -- so gradients still flow. QAT usually claws back the full FP32 accuracy at INT8 speed. The catch is obvious: it costs you a training run. So the sensible order of operations is to try post-training quantization first, measure the damage, and only reach for QAT if the damage is unacceptable. Don't pay for a retrain you didn't need.
Most networks are wildly overparameterized -- a huge fraction of the weights sit near zero and contribute almost nothing. Pruning finds and removes them:
import torch.nn.utils.prune as prune
model = SimpleClassifier()
# Unstructured pruning: zero out the smallest-magnitude weights.
# Here: drop the smallest 50% of weights in every Linear layer.
for name, module in model.named_modules():
if isinstance(module, nn.Linear):
prune.l1_unstructured(module, name='weight', amount=0.5)
def get_sparsity(model):
total, zeros = 0, 0
for p in model.parameters():
total += p.nelement()
zeros += (p == 0).sum().item()
return zeros / total
print(f"Sparsity: {get_sparsity(model):.1%}")
# Make pruning permanent (bake the zeros in, drop the mask)
for name, module in model.named_modules():
if isinstance(module, nn.Linear):
prune.remove(module, 'weight')
Here's the trap nobody tells beginners about, so I will: unstructured pruning makes your model sparse, but not automatically fast. You zeroed half the weights, sure, but the tensor is still exactly the same shape -- a 512x512 matrix full of zeros is still a 512x512 matrix as far as your CPU's matmul is concerned. To actually cash in that sparsity you need specialized sparse kernels or hardware that understands it. What does give you free speed on ordinary hardware is structured pruning, which removes whole neurons, channels, or attention heads, so the matrices genuinely get smaller:
# Structured pruning: remove entire output channels (neurons)
for name, module in model.named_modules():
if isinstance(module, nn.Linear):
prune.ln_structured(module, name='weight', amount=0.3, n=2, dim=0)
# drops 30% of output neurons, ranked by their L2 norm
The professional recipe is iterative pruning: cut a modest 10-20%, fine-tune for a few epochs to let the survivors compensate, measure accuracy, then repeat. Rip out 70% in one greedy pass and you'll wreck the model; sneak up on it in stages and you can often reach 50-80% sparsity with under 1% accuracy loss. Where exactly the wall sits depends entirely on your model and task -- there's no universal number, and anyone who quotes you one is selling something.
The techniques so far shrink a model you already have. Distillation flips the idea around: train a new, smaller model from scratch to imitate the big one.
class DistillationTrainer:
"""Train a small student to mimic a large teacher."""
def __init__(self, teacher, student, temperature=3.0, alpha=0.7):
self.teacher = teacher
self.student = student
self.temperature = temperature
self.alpha = alpha # weight on the soft (distillation) loss
self.teacher.eval()
self.optimizer = torch.optim.Adam(student.parameters(), lr=1e-3)
def distillation_loss(self, student_logits, teacher_logits, labels):
"""Blend soft targets from the teacher with the hard true labels."""
T = self.temperature
# Soft loss: KL divergence between softened teacher and student
soft_student = nn.functional.log_softmax(student_logits / T, dim=1)
soft_teacher = nn.functional.softmax(teacher_logits / T, dim=1)
soft_loss = nn.functional.kl_div(
soft_student, soft_teacher, reduction='batchmean'
) * (T * T) # scale by T^2 so gradient magnitudes stay sane
# Hard loss: ordinary cross-entropy against the real labels
hard_loss = nn.functional.cross_entropy(student_logits, labels)
return self.alpha * soft_loss + (1 - self.alpha) * hard_loss
def train_step(self, batch_x, batch_y):
with torch.no_grad():
teacher_logits = self.teacher(batch_x)
student_logits = self.student(batch_x)
loss = self.distillation_loss(student_logits, teacher_logits, batch_y)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item()
# Teacher: the big model (512 hidden, 3 layers)
teacher = SimpleClassifier(hidden=512)
# Student: a slim replacement (128 hidden, 2 layers)
student = nn.Sequential(
nn.Linear(784, 128), nn.ReLU(),
nn.Linear(128, 10),
)
trainer = DistillationTrainer(teacher, student, temperature=4.0, alpha=0.7)
The temperature is where the magic hides. Crank it up (3-5) and the teacher's probability distribution softens, exposing what Hinton beautifully called its dark knowledge -- the relative probabilities it quietly assigns to the wrong classes. A teacher that says "this handwritten 7 looks a little like a 1 and nothing at all like a 3" hands the student far richer information than a bare label saying "7". The student learns the shape of the teacher's uncertainty, not just its final answer. This is exactly how BERT (episode #59) begat DistilBERT: roughly 40% smaller, 60% faster, and still holding about 97% of the accuracy. Not a bad trade at all ;-)
ONNX (Open Neural Network Exchange) is an open, framework-agnostic format for models -- think of it as the PDF of neural networks. You export from PyTorch once, and the result runs on any runtime that speaks ONNX, on hardware PyTorch may not even support directly.
import torch.onnx
model = SimpleClassifier()
model.eval()
dummy_input = torch.randn(1, 784)
# Export the graph to ONNX by tracing one forward pass
torch.onnx.export(
model,
dummy_input,
"model.onnx",
input_names=['input'],
output_names=['output'],
dynamic_axes={'input': {0: 'batch_size'}}, # allow a variable batch size
opset_version=17,
)
# Run it with ONNX Runtime -- usually faster than PyTorch for inference
import onnxruntime as ort
import numpy as np
session = ort.InferenceSession("model.onnx")
input_data = np.random.randn(1, 784).astype(np.float32)
result = session.run(None, {'input': input_data})
print(f"Output shape: {result[0].shape}")
The lovely part is that ONNX Runtime is typically 1.5-3x faster than plain PyTorch for inference, with no other changes to your model. It earns that for free by fusing operations (folding a Linear-plus-ReLU into a single kernel, for instance), optimizing the memory layout, and leaning on CPU-specific instruction sets like AVX-512. That dynamic_axes argument is worth a second look too: without it, ONNX would bake in a fixed batch size of 1, and your carefully-exported model would refuse to process a batch of 32 in production.
If you're serving on NVIDIA hardware, TensorRT is the next rung up. It takes an ONNX model and compiles it into a highly-optimized inference engine tuned for your exact GPU -- it fuses layers, auto-selects the fastest kernel for each operation on that specific card, and can fold quantization right into the build.
# Typical path: PyTorch -> ONNX -> TensorRT engine.
# In practice you often build the engine straight from the ONNX file
# with the trtexec tool that ships with TensorRT:
#
# trtexec --onnx=model.onnx --saveEngine=model.plan --fp16
#
# The --fp16 (or --int8) flag tells TensorRT to build a reduced-precision
# engine, stacking quantization on top of every other optimization.
import tensorrt as trt
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(
1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)
parser = trt.OnnxParser(network, logger)
with open("model.onnx", "rb") as f:
parser.parse(f.read())
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16) # half precision on the GPU
engine = builder.build_serialized_network(network, config)
with open("model.plan", "wb") as f:
f.write(engine)
Because TensorRT optimizes against the actual card, an engine you build on an RTX 4090 is not the one you want running on an A100 in the cloud -- you build the engine on (or for) the hardware you'll deploy on. That's the small tax for the big reward: on GPU, a TensorRT engine with FP16 or INT8 can be several times faster than even ONNX Runtime, and it's a big part of how production inference servers hit the latencies they quote. A part from that, it composes cleanly with everything above -- distil, quantize, export to ONNX, then let TensorRT do the final squeeze.
Here's a hill I will die on: measuring model speed wrong is worse than not measuring at all, because a confident wrong number sends you optimizing the incorrect thing for a week.
def benchmark_model(model_fn, input_data, n_warmup=50, n_runs=200):
"""Measure inference latency without fooling yourself."""
# Warmup: the first runs are slow (JIT, cache warming, lazy init)
for _ in range(n_warmup):
model_fn(input_data)
# Make sure any queued GPU work is done before we start the clock
if torch.cuda.is_available():
torch.cuda.synchronize()
latencies = []
for _ in range(n_runs):
start = time.perf_counter()
model_fn(input_data)
if torch.cuda.is_available():
torch.cuda.synchronize() # WAIT for the GPU to actually finish
latencies.append(time.perf_counter() - start)
latencies = sorted(latencies)
print(f"Median latency: {np.median(latencies)*1000:.2f} ms")
print(f"P95 latency: {np.percentile(latencies, 95)*1000:.2f} ms")
print(f"P99 latency: {np.percentile(latencies, 99)*1000:.2f} ms")
print(f"Throughput: {1/np.median(latencies):.0f} inferences/sec")
# Compare FP32 vs quantized on the same input
model.eval()
test_input = torch.randn(1, 784)
with torch.no_grad():
benchmark_model(model, test_input)
benchmark_model(quantized_model, test_input)
Three rules turn that function from decoration into truth. One: always warm up. The first few calls pay for JIT compilation, cache misses, and lazy allocation -- include them and you're benchmarking the compiler, not the model. Two: synchronize the GPU. CUDA calls are asynchronous; they return the instant the work is queued, not when it's done. Skip torch.cuda.synchronize() and you'll clock the launch overhead and report a GPU that looks 100x faster than reality. Three: report percentiles, not the mean. Users don't feel your average latency; they feel the P99 -- the one request in a hundred that stalls. A model with a lovely 20ms mean and a hideous 800ms P99 will still get you angry emails. The mean hides exactly the tail that hurts, and the tail is where reputations go to die ;-)
None of these is the answer -- the real wins come from composing them. A fairly standard production pipeline looks like this:
Each step trades a little accuracy for a lot of speed, and because the trades are roughly independent, they multiply. Stack them and 10-50x faster inference at under 2% accuracy loss is a realistic outcome, not a fantasy. The discipline that makes it safe is the one we drilled last episode (#119): change one thing, measure, log it, move on. Optimize blind -- three tricks at once, no measurement -- and when accuracy tanks you'll have no idea which knob did it.
Three tasks to wrestle with before the next episode. As always, fight them yourself first -- I'll walk through full solutions next time.
Measure the real trade. Train SimpleClassifier on a toy dataset (MNIST, or even random data with fixed labels), then apply dynamic quantization. Use the benchmark_model function above to report median AND P99 latency for both the FP32 and INT8 versions, plus each model's size in MB and its accuracy. Write one sentence stating whether the speedup was worth the accuracy change for your case.
Sparse is not fast. Apply l1_unstructured pruning at 0%, 50%, and 90% sparsity to three copies of the model, then benchmark all three on CPU. You should find that latency barely moves even at 90% sparsity. Explain in a paragraph why the zeros don't speed anything up, and what would have to change (in hardware or kernels) for them to.
Distil, then measure the gap. Train a large teacher and a small student the ordinary way (student on hard labels only). Then train a second student of the same size using the DistillationTrainer. Score all three on a held-out set and report accuracy plus size. Bonus: sweep the temperature over {1, 2, 5, 10} and note how the distilled student's accuracy responds -- there's usually a sweet spot, and it's rarely 1.
That's the toolbox. A fast model, though, is still just sitting on your laptop being fast at nobody. In the next stretch we start wiring it into something the outside world can actually call -- and then worry about what changes when that same model has to run on a phone in your pocket, or survive months in production without anyone watching it. The engine is tuned; now we build the car around it.
@scipio