duan8/yolov5
2020-06-23 21:58:37 +08:00
..
CMakeLists.txt implement yolov5s 2020-06-22 13:48:17 +08:00
common.hpp implement yolov5s 2020-06-22 13:48:17 +08:00
gen_wts.py implement yolov5s 2020-06-22 13:48:17 +08:00
logging.h implement yolov5s 2020-06-22 13:48:17 +08:00
README.md update yolov5s 2020-06-23 21:58:37 +08:00
utils.h implement yolov5s 2020-06-22 13:48:17 +08:00
yololayer.cu yolov5s batchsize 2020-06-22 19:02:30 +08:00
yololayer.h yolov5s batchsize 2020-06-22 19:02:30 +08:00
yolov5s.cpp update yolov5s 2020-06-23 21:58:37 +08:00

yolov5

The Pytorch implementation is ultralytics/yolov5.

I was using [commit 6fd3939] of ultralytics/yolov5. And I made a copy of yolov5s.pt(google drive). Just in case the yolov5 model updated.

How to Run

1. generate yolov5s.wts from pytorch implementation with yolov5s.pt

git clone https://github.com/wang-xinyu/tensorrtx.git
git clone https://github.com/ultralytics/yolov5.git
// download its weights 'yolov5s.pt'
cd yolov5
cp ../tensorrtx/yolov5s/gen_wts.py .
python gen_wts.py
// a file 'yolov5s.wts' will be generated.

2. put yolov5s.wts into yolov5, build and run

mv yolov5s.wts ../tensorrtx/yolov5/
cd ../tensorrtx/yolov5
mkdir build
cd build
cmake ..
make
sudo ./yolov5s -s             // serialize model to plan file i.e. 'yolov5s.engine'
sudo ./yolov5s -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

Config

  • Input shape defined in yololayer.h
  • Number of classes defined in yololayer.h
  • FP16/FP32 can be selected by the macro in yolov5s.cpp
  • GPU id can be selected by the macro in yolov5s.cpp
  • NMS thresh in yolov5s.cpp
  • BBox confidence thresh in yolov5s.cpp
  • Batch size in yolov5s.cpp

More Information

See the readme in home page.