History

wang-xinyu d0a32e5646 yolov5 updated to v4.0		2021-01-08 16:52:48 +08:00
..
CMakeLists.txt	fix cmake	2020-12-16 12:54:00 +08:00
common.hpp	yolov5 updated to v4.0	2021-01-08 16:52:48 +08:00
gen_wts.py	update yolov5 genwts	2020-11-06 21:35:02 +08:00
logging.h	implement yolov5s	2020-06-22 13:48:17 +08:00
README.md	yolov5 updated to v4.0	2021-01-08 16:52:48 +08:00
utils.h	implement yolov5s	2020-06-22 13:48:17 +08:00
yololayer.cu	Fix bug : when the width and height are different, the result is incorrect (#235 )	2020-11-06 23:31:42 +08:00
yololayer.h	yolov5 v3.1 verified ok	2020-11-06 23:44:21 +08:00
yolov5_trt.py	fix: typo (#280 )	2020-11-11 09:42:00 +08:00
yolov5.cpp	yolov5 updated to v4.0	2021-01-08 16:52:48 +08:00

README.md

yolov5

The Pytorch implementation is ultralytics/yolov5.

Currently, we support yolov5 v1.0(yolov5s only), v2.0, v3.0, v3.1 and v4.0.

For yolov5 v4.0, please visit yolov5 release v4.0, and use the latest commit of this repo.
For yolov5 v3.0 and v3.1, please visit yolov5 release v3.0 and yolov5 release v3.1, and checkout commit '6d0f5cb' of this repo.
For yolov5 v2.0, please visit yolov5 release v2.0, and checkout commit '5cfa444' of this repo.
For yolov5 v1.0, please visit yolov5 release v1.0, and checkout commit 'f09aa3b' of this repo.

Config

Choose the model s/m/l/x by NET macro in yolov5.cpp
Input shape defined in yololayer.h
Number of classes defined in yololayer.h, DO NOT FORGET TO ADAPT THIS, If using your own model
FP16/FP32 can be selected by the macro in yolov5.cpp
GPU id can be selected by the macro in yolov5.cpp
NMS thresh in yolov5.cpp
BBox confidence thresh in yolov5.cpp
Batch size in yolov5.cpp

How to Run, yolov5s as example

1. generate yolov5s.wts from pytorch with yolov5s.pt, or download .wts from model zoo

git clone https://github.com/wang-xinyu/tensorrtx.git
git clone https://github.com/ultralytics/yolov5.git
// download its weights 'yolov5s.pt'
// copy tensorrtx/yolov5/gen_wts.py into ultralytics/yolov5
// ensure the file name is yolov5s.pt and yolov5s.wts in gen_wts.py
// go to ultralytics/yolov5
python gen_wts.py
// a file 'yolov5s.wts' will be generated.

2. build tensorrtx/yolov5 and run

// put yolov5s.wts into tensorrtx/yolov5
// go to tensorrtx/yolov5
// ensure the macro NET in yolov5.cpp is s
// update CLASS_NUM in yololayer.h if your model is trained on custom dataset
mkdir build
cd build
cmake ..
make
sudo ./yolov5 -s             // serialize model to plan file i.e. 'yolov5s.engine'
sudo ./yolov5 -d  ../samples // deserialize plan file and run inference, the images in samples will be processed.

3. check the images generated, as follows. _zidane.jpg and _bus.jpg

4. optional, load and run the tensorrt model in python

// install python-tensorrt, pycuda, etc.
// ensure the yolov5s.engine and libmyplugins.so have been built
python yolov5_trt.py

More Information

See the readme in home page.