YOLOv5-PyTorch

A PyTorch implementation of YOLOv5.

This repository has two features:

  • It is pure python code and can be run immediately using PyTorch 1.4 without build
  • Simplified construction and easy to understand how the model works

The model is based on ultralytics' repo,
and the code is using the structure of TorchVision.

Requirements

  • Windows or Linux, with Python ≥ 3.6

  • PyTorch ≥ 1.4.0

  • matplotlib - visualizing images and results

  • pycocotools - for COCO dataset and evaluation; Windows version is here

There is a problem with pycocotools for Windows. See Issue #356.

Besides, it's better to remove the prints in pycocotools.

optional:

  • nvidia dali (Linux) - a faster data loader

Datasets

This repository supports VOC and COCO datasets.

If you want to train your own dataset, you may:

  • write the correponding dataset code

  • convert your dataset to COCO-style

PASCAL VOC 2012 (download): http://host.robots.ox.ac.uk/pascal/VOC/voc2012/

MS COCO 2017: http://cocodataset.org/

Nvidia DALI is strongly recommended. It's much faster than the original data loader.

Currently this repository supports COCO-style dataset with DALI.

Training

Train on COCO dataset, using 1 GPU (if you wanna use 2 GPUs, set --nproc_per_node=2):

python -m torch.distributed.launch --nproc_per_node=1 --use_env train.py --use-cuda --dali --mosaic \
--epochs 190 --data-dir "./data/coco2017" --ckpt-path "yolov5s_coco.pth"

A more concrete modification is in run.sh.

To run it:

bash ./run.sh

If you are using PyTorch ≥ 1.6.0 and RTX series GPUs, the code will enable automatic mixed training (AMP).

Demo and Evaluation

  • Run demo.ipynb.

  • Modify the parameters in eval.ipynb to test the model.

Performance

Test on COCO 2017 val set, on a single RTX 2080Ti GPU:

The weights is from ultralytics' repo.

model bbox AP FPS params
YOLOv5s 36.1 410 7.5M

GitHub