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.