Transfer Learning for Vision
Teaching a giant a new task.
Orchard Vision has photographed fifteen hundred apples. Good ones and bruised ones. That's nowhere near enough to train a network from scratch, which would need millions. But it's plenty, if you borrow the eyes of a pretrained model.
- Fifteen hundred photos
- Far too few from scratch
- Borrow pretrained eyes
Features plus a new head
We remove the network's final layer, so it now gives five hundred and twelve numbers describing each image. We run all of our images through it once. Then we train a single linear layer on those numbers, which takes only seconds. We use small photos here in place of apples. Run it. Very strong accuracy, from fifteen hundred images, with almost no compute. This is often the first thing that you should try on a new vision problem.
import torch
from torch import nn
from torchvision import datasets, models
w = models.ResNet18_Weights.DEFAULT
tf = w.transforms()
train = torch.utils.data.Subset(datasets.CIFAR10("data", True, tf), range(1500))
test = torch.utils.data.Subset(datasets.CIFAR10("data", False, tf), range(500))
net = models.resnet18(weights=w).eval()
net.fc = nn.Identity()
def features(ds):
xs, ys = [], []
with torch.no_grad():
for x, y in torch.utils.data.DataLoader(ds, batch_size=64):
xs.append(net(x)); ys.append(y)
return torch.cat(xs), torch.cat(ys)
Xtr, ytr = features(train)
Xte, yte = features(test)
head = nn.Linear(512, 10)
opt = torch.optim.AdamW(head.parameters(), lr=1e-3)
for _ in range(150):
loss = nn.functional.cross_entropy(head(Xtr), ytr)
opt.zero_grad(); loss.backward(); opt.step()
print("Training images:", len(ytr))
print("Test accuracy:", round((head(Xte).argmax(1) == yte).float().mean().item(), 3))Training images: 1500 Test accuracy: 0.792
Getting more from your data
When data is tight, augmentation creates useful variations. Crops, flips, colour changes, and small rotations. Run it. One photo becomes eight quite different views. The label stays the same, so the model learns what really matters. For apples, colour jitter should be gentle though, since colour tells you about ripeness. Always ask what your augmentation might break.
import torch
import numpy as np
from PIL import Image
from torchvision.transforms import v2
from skimage import data
import matplotlib.pyplot as plt
image = Image.fromarray(data.chelsea())
aug = v2.Compose([v2.RandomResizedCrop(240, scale=(0.5, 1.0)), v2.RandomHorizontalFlip(),
v2.ColorJitter(0.4, 0.4, 0.4, 0.05), v2.RandomRotation(12)])
torch.manual_seed(0)
fig, ax = plt.subplots(2, 4, figsize=(9, 4.4))
for a in ax.ravel():
a.imshow(aug(image)); a.axis("off")
plt.tight_layout(); plt.savefig("augment.png", dpi=110)
print("One apple photo can become endless training examples")One apple photo can become endless training examples
Which mistakes cost more?
Finally, measure the right thing. Here are one hundred apples, ten of them bruised. Accuracy is ninety five percent. But look at bruised recall in the report. It's only seventy percent, so three bad apples slipped through. In a food business, that matters more than the overall number. Always read the rare class.
import numpy as np
from sklearn.metrics import confusion_matrix, classification_report
truth = np.array(["good"] * 90 + ["bruised"] * 10)
pred = truth.copy()
pred[np.array([3, 17])] = "bruised"
pred[np.array([92, 95, 98])] = "good"
print(confusion_matrix(truth, pred, labels=["good", "bruised"]))
print(classification_report(truth, pred, digits=2))[[88 2]
[ 3 7]]
precision recall f1-score support
bruised 0.78 0.70 0.74 10
good 0.97 0.98 0.97 90
accuracy 0.95 100
macro avg 0.87 0.84 0.85 100
weighted avg 0.95 0.95 0.95 100Putting a price on mistakes
A missed bruise might cost two dollars in returns. Throwing away a good apple costs twenty five cents. So where's the best setting? We compare three thresholds, in dollars. Run it. The cheapest option isn't always the most accurate one. Now the owner can make a business decision. This is how a vision system goes from a nice demo to something that pays for itself.
cost_missed_bruise = 2.0
cost_false_reject = 0.25
apples = 100000
rate = 0.04
for name, miss, reject in [("Strict", 0.05, 0.08), ("Balanced", 0.15, 0.03), ("Loose", 0.40, 0.01)]:
missed = apples * rate * miss * cost_missed_bruise
wasted = apples * (1 - rate) * reject * cost_false_reject
print(f"{name:<9} missed ${missed:>6,.0f} wasted ${wasted:>6,.0f} total ${missed + wasted:>6,.0f}")Strict missed $ 400 wasted $ 1,920 total $ 2,320 Balanced missed $ 1,200 wasted $ 720 total $ 1,920 Loose missed $ 3,200 wasted $ 240 total $ 3,440
Summary
Pretrained features, gentle augmentation, and the right metric. AppleEye can now grade an apple in a photo. But the sorting line has many apples, all at once. Next, finding each object in a scene.
Key takeaways- Replace the head, train it
- Augment with care
- Price your mistakes