YOLO Object Detection Model

YOLO (“You Only Look Once”) was the paper that made object detection fast: one pass through one network, and out come all the boxes. I’d trained classifiers by then, and a classifier only has to say what is in a picture. A detector has to say what, where and how many, all at once. I wanted to know how a single network pulls that off, so I built the original 2016 version from the paper.

TLDR: YOLOv1 in PyTorch: a 24-layer convolutional backbone, a 7×7 grid, two boxes per cell, and the paper’s sum-of-squares loss. My first full run (Aug 4) found nothing in five of six test images and drew twenty junk boxes on the sixth. Three days and a lot of debugging later, it labels zebras, giraffes, clocks and trains correctly. It’s trained on COCO128, a 128-image sample of COCO, so think of it as proof that the pipeline works end to end.

How YOLO sees an image

The image gets cut into a 7×7 grid. Each object belongs to exactly one cell: the one its centre falls in. That cell has to predict the box (as an offset inside the cell, plus a width and height relative to the whole image) and the class.

Street scene with a 7 by 7 grid; boxes around a clock, a stop sign, a truck and a car, and the cell containing each box's centre shaded — Drawn from the image's label file: four objects, four centre points, and the four cells that are responsible for them. Every other cell should predict 'nothing here'.

That rule has a funny side effect. If two objects have their centres in the same cell, YOLOv1 can only learn one of them. The code just skips the second.

Python source that builds the 7 by 7 training target, one box per cell — Building the target. Note the 'cell already taken' line: that's a real limit of YOLOv1, not a shortcut.

The loss

The loss is where most of the paper lives. Each cell predicts two boxes, but only the one that overlaps the real object best is “responsible” for it. That box gets pushed toward the right position and size; the other one only gets told it should be less confident. Widths and heights go through a square root first, so being 10 pixels off on a small box hurts more than on a big one.

Python source for the YOLO loss: picking the responsible box by IoU, then box, object, no-object and class terms — The whole loss. Box error is weighted ×5 and empty cells ×0.5, straight from the paper.

The first run

On August 4 I trained it, ran it on six images, and got this.

Six photos: five with zero detections and one covered in twenty overlapping boxes with wrong labels — The failed run. Five images with nothing. The sixth has twenty boxes, labelled things like elephant, motorcycle and skateboard. It's a dog on some shoes.

Nothing looks wrong when you read the code, but it took five separate fixes to get it working.

Not enough training. 128 images at batch size 8 is 16 steps per epoch. Fifty epochs is 800 updates, for a network with 284 million weights. I tripled it.

The head was most of the network. The paper’s final fully connected layer maps 50,176 features to 4,096. That one layer is 224M of the 284M weights: 80% of the model trying to memorise 128 pictures. Shrinking it to 1,024 and dropping the dropout made it converge.

Negative boxes. The network’s outputs aren’t constrained, so early on it happily predicts boxes with negative width. Non-maximum suppression, the step that removes duplicate boxes, treats those as having negative area and never removes them. That’s where the twenty-box pile came from. Clamping widths at zero, and dropping empty boxes, fixed it.

A square root at zero. The square root on box size has an infinite slope at zero. Any width that drifted through zero sent a huge gradient straight into the head. A tiny epsilon under the root, plus gradient clipping, stopped the blowups.

BatchNorm lying about its stats. This was the sneaky one. The model scored a loss of 15.9 in training mode and 24.2 in evaluation mode, same weights, same data. BatchNorm keeps a running average of its statistics during training, and at batch size 8 that average lags far behind the weights. So training looked fine and predictions were garbage. The fix is one extra pass over the data at the end to recompute the averages exactly.

Python source for recalibrate_batchnorm, which resets each BatchNorm's running statistics and recomputes them in one pass — Recomputing BatchNorm's running stats after training. The docstring has the 15.9 vs 24.2 that gave it away.

The working run

Same six images (plus four more), August 7.

The same photo before and after: twenty wrong boxes on the left, a single box labelled dog 0.81 on the right — Same image, three days apart. Twenty wrong boxes, then one right one.

Ten photos with correct detections: bowls, oranges and broccoli, giraffes, a vase, a zebra, a person with an umbrella, a dog, horses, a clock and a train — The final run. It's seen these images in training, so this shows the pipeline works. Generalising to new images needs much more data.

What’s next

Train on real data. With a pretrained backbone and full COCO instead of 128 images, the bigger head and the dropout from the paper should come back. YOLOv2’s anchor boxes are the fix for the one-object-per-cell limit, so that’s the natural next step.