CSE445Computer Vision
Tutorials/Object Detection/YOLO Training on a Custom Dataset
Previous tutorialYOLO Architecture: Backbone, Neck & HeadNext tutorialOD 8A: DETR Transformer Detection
OD-07Intermediate 150 min

YOLO Training on a Custom Dataset

End-to-end workflow: annotate images in YOLO format, configure a training run, monitor loss curves, evaluate with COCO metrics, and tune augmentations for your domain.

Ultralytics YOLOv8RoboflowTensorBoardONNX Runtime

Code Structure Note

These notes are provided only to help students understand the code structure. Students should review, optimize, and adapt the code before using it in assignments, projects, or experiments.

Topics Covered
  1. 1Dataset annotation with LabelImg / Roboflow; YOLO label format (class cx cy w h)
  2. 2data.yaml: paths, nc, class names — the single source of truth for a training run
  3. 3Transfer learning: freeze backbone layers for the first N epochs on small datasets
  4. 4Augmentation pipeline: mosaic, mixup, random flip, HSV jitter, copy-paste
  5. 5Training curves: box loss, cls loss, dfl loss, mAP@50, mAP@50-95

Learning Objectives

  • Annotate a custom dataset and export it in YOLO format using LabelImg or Roboflow
  • Write a valid data.yaml and launch a YOLOv8 training run from the CLI
  • Interpret the TensorBoard training curves and identify under/overfitting
  • Tune mosaic probability and freeze epochs and measure the effect on val mAP

Key Takeaways

  • Freezing backbone layers for the first 10 epochs prevents early over-writing of ImageNet features
  • Mosaic augmentation tiles 4 images into one training sample — dramatically improves small-object recall
  • 80–20 train/val split is the minimum; class balance matters more than raw image count
  • A 50-epoch run on 500 images takes <10 min on a T4 GPU — iterate fast, measure on val

Tutorial Infographic

Notebook Resources

Jupyter Notebook

Student action area

Use this section to practice the tutorial code. Kaggle opens the prepared notebook, Colab lets you experiment in your own runtime, and the .ipynb download is best for local editing or submission.

Open in Kaggle Open in Colab Download .ipynb
Kaggle: follow the official tutorial notebook.
Colab: modify and rerun quickly.
IPYNB: keep a local copy.

Hands-On Exercises

  1. 1Annotate 50 images of a custom object using LabelImg and export to YOLO format
  2. 2Train YOLOv8s for 50 epochs; plot box loss and mAP@50 curves from `results.csv`
  3. 3Compare mAP@50 with freeze=10 vs. freeze=0 on the same dataset
  4. 4Export the trained model to ONNX and run the inference script on 5 held-out test images

Tools & Libraries

Ultralytics YOLOv8RoboflowTensorBoardONNX Runtime
Previous tutorialYOLO Architecture: Backbone, Neck & HeadNext tutorialOD 8A: DETR Transformer Detection