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.ipynbto test the model.
Performance
Test on COCO 2017 val set, on a single RTX 2080Ti GPU:
The weights is from ultralytics’ repo.
modelbbox APFPSparamsYOLOv5s36.14107.5M
#machine learning #pytorch #yolov5 #python
