YOLO Object Detector
YOLOv1 from the Paper, in PyTorch
Give it an image and it draws a box around every object it knows, with its class, in one pass.
- mAP@0.5 0.952
- 8.3 M parameters, CPU in ~25 min
- Paper network: 271.7 M parameters
Built with
PyTorchtorchvision
pytestMatplotlib
NumPy
Kaggle
Google Colab
Overview
YOLOv1 turns object detection into one regression: the image is split into a 7 × 7 grid, and the cell that holds an object's centre predicts its class and two candidate boxes with a confidence. The network sees the whole image at once, so detection is a single forward pass. I wrote every part of it in PyTorch from the paper: the 24-layer network, the sum-squared-error loss, target encoding, decoding, non-max suppression and mAP, with no detection library and no pretrained weights.
The project started in 2023 as two notebooks trained on Mapillary traffic signs. When I rewrote it as a tested package I found that the test and validation CSVs were subsets of the training CSV, so every mAP had been measured on training images; that classes above 20 were never counted; that boxes were drawn from their centre as if it were a corner; and that the second box of each cell was never used at test time. The legacy folder keeps the notebooks and the full list of fixes.
To show the pipeline learns without ImageNet pretraining I added two datasets that run on a CPU: synthetic shapes and MNIST digits scattered on a canvas. A small 8.3M-parameter network reaches validation mAP@0.5 of 0.952 on the shapes and 0.921 on the digits in about 25 minutes each. PASCAL VOC and any YOLO-format folder are supported through the same command and a Kaggle / Colab notebook, but those GPU runs are not reported.
The question
Can a single-pass detector written without detection libraries or pretrained weights learn to find and classify objects, and how do I prove that the loss, decoding and mAP code is right?
Key features
- Two architectures behind one flag: the paper's 24 conv + 2 FC network (271.7M parameters, 448 px) and a 6-conv small network (8.3M parameters, 112 px) for CPU runs, both with batch norm and a width multiplier.
- The paper loss for any number of boxes per cell: the best-IoU predictor is responsible, λ_coord = 5 and λ_noobj = 0.5, square-root box sizes, and the four components logged every epoch.
- Decoding with the class-specific score confidence × P(class), per-class non-max suppression and VOC-style mAP@0.5 with all-point interpolation.
- Four datasets through one flag: on-the-fly shapes, MNIST digits on a canvas, PASCAL VOC 2007 / 2012, and any images/ + labels/ folder in the YOLO text format with a seeded train / validation split.
- Paper-style augmentation (scale, translation, flip, brightness and saturation jitter) that moves and clips the boxes, with a switch to turn it off.
- Three commands: yolov1-train writes best.pt, history curves, a prediction grid and per-class AP; yolov1-evaluate reports mAP; yolov1-predict draws boxes on files or a webcam feed.
- A pytest suite that checks IoU, encoding and decoding, each loss component, mAP and NMS against hand-computed values, plus an overfit sanity check on four images.
How it works
Encode the targets
Each ground-truth box goes to the grid cell that contains its centre, as an offset inside the cell, a size relative to the image and a one-hot class, giving an S × S × (C + 5B) tensor.
One forward pass
Conv, batch-norm and leaky-ReLU blocks with four max-pools reduce the image to a 7 × 7 map; two fully connected layers output every box, confidence and class probability at once.
Sum-squared-error loss
For every object the box with the highest IoU is responsible and is pulled to the target coordinates and to confidence 1; every other predictor is pushed to confidence 0, weighted by λ_noobj; class probabilities are regressed in object cells.
Decode and suppress
Every box gets the score confidence × best class probability; boxes under the threshold are dropped and same-class boxes that overlap a kept box above the NMS IoU are removed.
Measure mAP
Per class, detections are sorted by score, matched to unmatched ground truth at IoU ≥ 0.5, and AP is the area under the precision envelope; mAP averages the classes that have ground truth.
Examples
Input
yolov1-train --dataset mnist --root data --image-size 112 --arch small --hidden 512 --dropout 0.3 --no-augment --train-size 10000 --val-size 500 --epochs 12 --batch-size 32 --lr 1e-3 --output-dir checkpoints/mnist
Output
checkpoints/mnist/: best.pt (highest validation mAP), last.pt, history.json, history.png, predictions.png, best_metrics.json — validation mAP@0.5 0.921 after 12 epochs, about 26 minutes on an 8-core CPU.
Tech stack
The network, the loss, decoding, NMS, mAP and the training loop
MNIST and PASCAL VOC downloads and image transforms
Hand-computed checks of boxes, loss, metrics, data and an overfit sanity test
Training curves and prediction grids
Shape generation and box bookkeeping
GPU training notebook for VOC and YOLO-format datasets
The same notebook, generated from src by build_notebook.py
What I learned
- A detector that cannot memorise four images has a bug in its data, loss or decoding, so the overfit test is the most useful test in the suite
- Augmentation has to respect the labels: a horizontally flipped digit is a different digit, so the MNIST run trains without flips
Limitations
- No PASCAL VOC or Mapillary run is reported: the paper network needs a GPU, and the published numbers are small CPU runs on simple data
- The 2023 notebooks this started from measured mAP on training images, so none of their numbers are comparable with the package
Keywords
YOLO · YOLOv1 · object detection · You Only Look Once · PyTorch · non-maximum suppression · NMS · mean average precision · mAP · PASCAL VOC · MNIST · Mapillary traffic signs · computer vision · deep learning