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.
-
Windows or Linux, with Python ≥ 3.6
-
PyTorch ≥ 1.6.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(recommended). Download page
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: http://host.robots.ox.ac.uk/pascal/VOC/voc2012/
MS COCO 2017: http://cocodataset.org/
NVIDIA DALI is strongly recommended. It may be much faster than the original data loader.
Currently this repository supports COCO-style dataset with DALI.
The model is mainly made of Darknet and PANet.
You can get its flowchart by opening images/YOLOv5.drawio
with drawio.
Train on COCO dataset, using 1 GPU (if you wanna use N GPUs, just set --nproc_per_node=N):
python -m torch.distributed.run --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 use RTX series GPUs, the code will enable automatic mixed training (AMP).
- Run
demo.ipynb
.
- Modify the parameters in
eval.ipynb
to test the model.
Test on COCO 2017 val set, on 1 2080Ti GPU:
The weights is from ultralytics' repo.
model | bbox AP | FPS | params |
---|---|---|---|
YOLOv5s | 36.1 | 410 | 7.5M |