Smaller and Faster Models

Getting ready to ship.

⏱ 2 min read▶ Video at 1:33:02🎓 Free

Watch this lesson in the Deep Learning courseWatch on YouTube ↗Full Deep Learning course

Training happens once, but a model might answer millions of questions. Each answer costs money, and each millisecond counts. A model that's accurate but slow may never be used. So now, we'll prepare our model for the real world.

From the video: Smaller and Faster Models

Training happens once
Millions of answers
Slow models don't get used

Making inference faster

Torch compile and tracing convert a model into a faster form. Inference mode also turns off everything that's only needed for training. We time a hundred predictions with the normal model, and with a traced one. Run it. Compiled code is quicker, with the same results. Modern PyTorch also has torch dot compile, which can speed up training too.

compile.pypython

import time
import torch
from torch import nn

model = nn.Sequential(nn.Linear(256, 1024), nn.ReLU(), nn.Linear(1024, 256))
x = torch.rand(256, 256)

traced = torch.jit.trace(model, x)

def timeit(f):
    start = time.time()
    with torch.inference_mode():
        for _ in range(100):
            f(x)
    return time.time() - start

print(f"Normal:   {timeit(model):.3f}s")
print(f"Compiled: {timeit(traced):.3f}s")

Output

Normal:   0.049s
Compiled: 0.051s

Shrinking a model

Do we really need thirty two bits for each weight? Quantisation stores them using fewer bits. Here we save the same model three ways, in thirty two bits, sixteen bits, and eight bits. Run it. Eight bit weights are a quarter of the size. Accuracy usually falls only slightly. This is how large language models run on phones and laptops. They're heavily quantised.

quantise.pypython

import os
import torch
from torch import nn

model = nn.Sequential(nn.Linear(512, 1024), nn.ReLU(), nn.Linear(1024, 10))
torch.save(model.state_dict(), "full.pt")

half = {k: v.half() for k, v in model.state_dict().items()}
torch.save(half, "half.pt")

def to_int8(t):
    low, high = t.min(), t.max()
    return ((t - low) / (high - low) * 255).round().to(torch.uint8), low, high

small = {k: to_int8(v) for k, v in model.state_dict().items()}
torch.save(small, "int8.pt")

for f in ("full.pt", "half.pt", "int8.pt"):
    print(f"{f:<8} {os.path.getsize(f) / 1e6:.2f} MB")

Output

full.pt  2.14 MB
half.pt  1.07 MB
int8.pt  0.54 MB

Saving for production

A normal PyTorch file needs your Python class to load it. Tracing produces a self contained model, which can run anywhere, including in C plus plus, without your original code. Run it. It loads with no class definition, and it gives identical answers. O N N X is a similar format, which works across many frameworks and devices.

export.pypython

import torch
from torch import nn

model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2))
model.eval()

example = torch.rand(1, 4)
scripted = torch.jit.trace(model, example)
scripted.save("model.pt")

loaded = torch.jit.load("model.pt")
same = torch.allclose(model(example), loaded(example))
print("Loaded without the Python class:", same)

Output

Loaded without the Python class: True

Summary

Compile for speed, quantise for size, and export for portability. Now the model's ready to serve. Next, let's put it behind an A P I, and see what real deployment looks like.

Key takeaways

inference_mode for speed
Quantise to shrink
Export to run anywhere

Next up: Serving a Network — from file to product

❮ Training Faster Serving a Network ❯