From a Function to a Service
The interface, where preprocessing belongs, and loading once
A trained model is a function from tensors to tensors. A product is something a caller can send a request to and get an answer from, that keeps working when the caller sends something unexpected, and that tells you when it has stopped being right. This week builds the distance between those two things.
The interface is the design decision
Before any code, decide what crosses the boundary. The model wants a normalised float tensor of a particular shape. A caller has a picture. Somewhere between the two sits preprocessing, and where you put it decides what can go wrong.
- Caller sends raw input, service preprocesses. One implementation of the preprocessing, in one place, versioned with the model. Almost always right.
- Caller sends a preprocessed tensor. Now every caller has its own copy of your normalisation constants, and they will drift apart. Day 3 measures what that costs.
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets, transforms
MEAN, STD = 0.1307, 0.3081
tf_ = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((MEAN,), (STD,))])
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 build():
return nn.Sequential(
nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
nn.Flatten(), nn.Linear(32 * 7 * 7, 64), nn.ReLU(),
nn.Linear(64, 10))
def get_model(path='served.pt', epochs=6):
"""Train once, then reuse the checkpoint on later runs."""
model = build()
if os.path.exists(path):
model.load_state_dict(torch.load(path))
model.eval()
return model
torch.manual_seed(0)
model = build()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(xtr, ytr), batch_size=128,
shuffle=True)
for _ in range(epochs):
model.train()
for xb, yb in loader:
opt.zero_grad()
lf(model(xb), yb).backward()
opt.step()
model.eval()
torch.save(model.state_dict(), path)
return model
def accuracy(model, x, y):
model.eval()
with torch.no_grad():
return (model(x).argmax(1) == y).float().mean().item()
from fastapi import FastAPI, HTTPException
from fastapi.testclient import TestClient
from pydantic import BaseModel, Field
model = get_model()
app = FastAPI()
class Request(BaseModel):
# 784 raw pixels in 0 to 255, exactly what a caller would send
pixels: list[float] = Field(min_length=784, max_length=784)
def preprocess(pixels):
"""The one function that must match training. Note the normalise:
it is the same MEAN and STD the training transform used, and it is
the single most common thing to get wrong here."""
x = torch.tensor(pixels, dtype=torch.float32).reshape(1, 1, 28, 28)
return (x / 255.0 - MEAN) / STD
@app.post('/predict')
def predict(req: Request):
if max(req.pixels) > 255 or min(req.pixels) < 0:
raise HTTPException(422, 'pixels must be between 0 and 255')
x = preprocess(req.pixels)
with torch.no_grad():
probs = torch.softmax(model(x), dim=1)[0]
top = int(probs.argmax())
return {'digit': top, 'confidence': round(float(probs[top]), 4)}
@app.get('/health')
def health():
return {'status': 'ok', 'params': sum(p.numel()
for p in model.parameters())}
client = TestClient(app)
def raw_pixels(i):
"""Undo the training normalisation to get back what a caller sends."""
return ((xte[i, 0] * STD + MEAN) * 255.0).flatten().tolist()
print(client.get('/health').json())
r = client.post('/predict', json={'pixels': raw_pixels(0)})
print('status %d %s' % (r.status_code, r.json()))
print('true label %d' % yte[0])
status 200 {'digit': 7, 'confidence': 1.0}
true label 7
That is a complete service. The request carries 784 raw pixel values in the range a caller would actually have, the service normalises them using the same constants the training transform used, and it returns a digit and a confidence. The health endpoint exists so that whatever is running the service can tell whether it came up.
Return the confidence, always
It costs nothing, and without it the caller cannot distinguish a prediction the model is certain of from one it is guessing at. It is also the number day 6 uses to detect that the model has stopped working, and you cannot collect it retrospectively.
Why the model loads once
The model is constructed at import time and held in memory. Loading a checkpoint takes far longer than a forward pass, so a service that loads per request spends nearly all of its time on the wrong thing. This is stated because it is a genuinely common mistake, and because it is invisible until you measure.