Project: A Simple Object Detector

Build the mental model

Object detection can be simplified down to a regression problem: instead of predicting one class label per image, you predict 4 continuous numbers — the x1, y1, x2, y2 corners of a bounding box. The backbone is the same Conv2d/ReLU stack from earlier projects extracting spatial features, but here it feeds into AdaptiveAvgPool2d(1) for global average pooling, producing a fixed-length vector regardless of input resolution. The head is a small Linear layer followed by Sigmoid, squashing outputs into the [0, 1] range to match the standard convention of representing box coordinates as fractions of image width and height.

For training, SmoothL1Loss (a Huber-style loss) measures the regression error between predicted and ground-truth boxes — it's more robust to occasional large coordinate errors than plain MSE. But it's worth noticing the mismatch between the training loss and the evaluation metric that actually matters: Intersection over Union (IoU) is what you'd use at test time to judge detection quality, yet it isn't smooth or differentiable enough near zero overlap to backpropagate through directly, which is exactly why a smooth surrogate loss like SmoothL1 is used instead. Reusing the from-scratch IoU function built earlier in the course to sanity-check that a falling regression loss actually corresponds to improved box overlap is the core idea this project is built around.

Connect it to a real scenario

The Tutorial Platform's thumbnail auto-crop tool uses exactly this kind of simplified detector to avoid making instructors manually pick a crop region: it predicts a bounding box around the main subject (a code window or a presenter's face) in an uploaded course thumbnail, and that predicted box is used directly as the crop region for catalog cards.

Object detection can be simplified into a 4-number bounding-box regression problem, with a Sigmoid head constraining outputs to the [0,1] normalized-coordinate range.
SmoothL1Loss serves as a differentiable training loss, but it's distinct from IoU, the metric that actually matters for detection quality and isn't suitable for direct backpropagation.
A from-scratch IoU function is a reusable sanity check for confirming that a falling regression loss actually corresponds to improved box overlap.

Try the working example

python

import torch
import torch.nn as nn
import torch.optim as optim

torch.manual_seed(0)

def box_iou(box_a, box_b):
    # box_a, box_b: tensors of shape (4,) as [x1, y1, x2, y2] in normalized coords
    x1 = torch.max(box_a[0], box_b[0])
    y1 = torch.max(box_a[1], box_b[1])
    x2 = torch.min(box_a[2], box_b[2])
    y2 = torch.min(box_a[3], box_b[3])

    inter_w = (x2 - x1).clamp(min=0)
    inter_h = (y2 - y1).clamp(min=0)
    intersection = inter_w * inter_h

    area_a = (box_a[2] - box_a[0]).clamp(min=0) * (box_a[3] - box_a[1]).clamp(min=0)
    area_b = (box_b[2] - box_b[0]).clamp(min=0) * (box_b[3] - box_b[1]).clamp(min=0)
    union = area_a + area_b - intersection

    return (intersection / union).item() if union > 0 else 0.0

class TinyDetector(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),  # 32x32 -> 16x16
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),  # -> 32x1x1, works for any input size
        )
        self.head = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32, 4),
            nn.Sigmoid(),  # squash to [0, 1] like normalized box coords
        )

    def forward(self, x):
        x = self.features(x)
        return self.head(x)

num_samples = 32
images = torch.randn(num_samples, 3, 32, 32)

# Synthetic ground-truth boxes: random but valid [x1, y1, x2, y2], x2>x1, y2>y1
x1 = torch.rand(num_samples, 1) * 0.5
y1 = torch.rand(num_samples, 1) * 0.5
x2 = x1 + torch.rand(num_samples, 1) * 0.5 + 0.1
y2 = y1 + torch.rand(num_samples, 1) * 0.5 + 0.1
gt_boxes = torch.cat([x1, y1, x2, y2], dim=1).clamp(max=1.0)

model = TinyDetector()
criterion = nn.SmoothL1Loss()
optimizer = optim.Adam(model.parameters(), lr=1e-2)

model.eval()
with torch.no_grad():
    pred_before = model(images)[0]
iou_before = box_iou(pred_before, gt_boxes[0])
print(f"IoU before training (sample 0): {iou_before:.4f}")

model.train()
for step in range(20):
    optimizer.zero_grad()
    preds = model(images)
    loss = criterion(preds, gt_boxes)
    loss.backward()
    optimizer.step()
    if step % 5 == 0 or step == 19:
        print(f"Step {step+1}/20 - loss: {loss.item():.4f}")

model.eval()
with torch.no_grad():
    pred_after = model(images)[0]
iou_after = box_iou(pred_after, gt_boxes[0])
print(f"IoU after training (sample 0): {iou_after:.4f}")
print("Training complete.")

You should see

Output starts with "IoU before training (sample 0): X.XXXX" — since the untrained head's output sits near the sigmoid midpoint (0.5) in each of the 4 coordinates, overlap with the ground-truth box tends to be small (IoU likely under 0.3, possibly 0.0 if there's no overlap at all). Then five loss lines print for steps 1, 6, 11, 16, and 20 — with only 32 samples to fit via regression, SmoothL1 loss should decrease clearly over the 20 steps. Finally, "IoU after training (sample 0): X.XXXX" should typically be higher than the before value, followed by "Training complete."

5-minute try-it

Extend the IoU check to compute the mean IoU across all 32 samples, before and after training, instead of just sample 0. Compare the mean IoU improvement against a version that swaps in criterion = nn.MSELoss() instead of SmoothL1Loss.

One important caution

Generating ground-truth boxes without preserving the x2>x1 and y2>y1 constraints can produce degenerate boxes with zero or negative width/height, which the clamp(min=0) logic inside box_iou will silently treat as zero area rather than raising an error.

Assuming a falling SmoothL1 loss automatically means rising IoU is a trap — they correlate on average but aren't the same signal, so checking IoU on just one sample (as this project does for sample 0) can show a dip even while the overall loss trends down.

PyTorch Vision Docs — Models and pre-trained weights — Computer Vision