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 only) - 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 PyTorch's 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.