A Small Bruise Detector
A network you fully understand.
Transfer learning is great, but it hides a lot. So let's build a tiny bruise detector, with nothing borrowed. It's small enough to train in a minute on a laptop. And in the next few lessons, we'll break it on purpose, in all the ways that real vision systems break.
- Nothing borrowed
- Trains in a minute
- We'll break it on purpose
Patches and a tiny network
This helper file does three things. Make draws small simulated apple photos, and half of them get a bruise. Net is a tiny network, with two convolution layers. Train is the normal loop. Accuracy scores a model. We'll use this file for the rest of the module. Notice that make has a light setting, because we're going to need that very soon.
import numpy as np
import cv2
import torch
def make(n, seed=0, light=1.0, bruise_rate=0.5):
"""Simulated apple photos. Label 1 means bruised."""
rng = np.random.default_rng(seed)
X, y = [], []
for _ in range(n):
img = np.full((48, 48, 3), (70, 75, 80), np.uint8)
c = tuple(int(v) for v in rng.integers(19, 30, 2))
r = int(rng.integers(15, 20))
cv2.circle(img, c, r, (35, 40, int(rng.integers(150, 225))), -1, cv2.LINE_AA)
bad = rng.random() < bruise_rate
if bad:
at = (c[0] + int(rng.integers(-8, 9)), c[1] + int(rng.integers(-8, 9)))
cv2.circle(img, at, int(rng.integers(3, 6)), (25, 45, 70), -1, cv2.LINE_AA)
img = img * light + rng.normal(0, 8, img.shape)
X.append(img.clip(0, 255) / 255)
y.append(int(bad))
X = torch.tensor(np.array(X), dtype=torch.float32).permute(0, 3, 1, 2)
return X, torch.tensor(y)
def net():
return torch.nn.Sequential(
torch.nn.Conv2d(3, 16, 3, padding=1), torch.nn.ReLU(), torch.nn.MaxPool2d(2),
torch.nn.Conv2d(16, 32, 3, padding=1), torch.nn.ReLU(), torch.nn.MaxPool2d(2),
torch.nn.Flatten(), torch.nn.Linear(32 * 12 * 12, 2))
def train(model, X, y, epochs=8):
opt = torch.optim.Adam(model.parameters(), lr=2e-3)
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 accuracy(model, X, y):
return (model(X).argmax(1) == y).float().mean().item()Training it
We make fifteen hundred training images, and five hundred new ones for testing. Run it. Just fourteen thousand weights, and it's nearly ninety eight percent accurate. That's a tiny model, in one minute. It looks like a finished product. Most tutorials stop right here. But production begins here. So let's find out what's really going on.
import torch
from patch import make, net, train, accuracy
torch.manual_seed(0)
X, y = make(1500)
X_te, y_te = make(500, seed=9)
model = train(net(), X, y)
print("Weights:", sum(p.numel() for p in model.parameters()))
print(f"Test accuracy: {accuracy(model, X_te, y_te):.3f}")
torch.save(model.state_dict(), "bruise.pt")Weights: 14306 Test accuracy: 0.976
Also uses patch.py
import numpy as np
import cv2
import torch
def make(n, seed=0, light=1.0, bruise_rate=0.5):
"""Simulated apple photos. Label 1 means bruised."""
rng = np.random.default_rng(seed)
X, y = [], []
for _ in range(n):
img = np.full((48, 48, 3), (70, 75, 80), np.uint8)
c = tuple(int(v) for v in rng.integers(19, 30, 2))
r = int(rng.integers(15, 20))
cv2.circle(img, c, r, (35, 40, int(rng.integers(150, 225))), -1, cv2.LINE_AA)
bad = rng.random() < bruise_rate
if bad:
at = (c[0] + int(rng.integers(-8, 9)), c[1] + int(rng.integers(-8, 9)))
cv2.circle(img, at, int(rng.integers(3, 6)), (25, 45, 70), -1, cv2.LINE_AA)
img = img * light + rng.normal(0, 8, img.shape)
X.append(img.clip(0, 255) / 255)
y.append(int(bad))
X = torch.tensor(np.array(X), dtype=torch.float32).permute(0, 3, 1, 2)
return X, torch.tensor(y)
def net():
return torch.nn.Sequential(
torch.nn.Conv2d(3, 16, 3, padding=1), torch.nn.ReLU(), torch.nn.MaxPool2d(2),
torch.nn.Conv2d(16, 32, 3, padding=1), torch.nn.ReLU(), torch.nn.MaxPool2d(2),
torch.nn.Flatten(), torch.nn.Linear(32 * 12 * 12, 2))
def train(model, X, y, epochs=8):
opt = torch.optim.Adam(model.parameters(), lr=2e-3)
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 accuracy(model, X, y):
return (model(X).argmax(1) == y).float().mean().item()Summary
A tiny network can look brilliant. Next, how to choose where to draw the line.
Key takeaways- Small networks train fast
- Test on new data
- High accuracy isn't the end