The Pytorch implementation is ultralytics/yolov5.
Currently, we support yolov5 v1.0(yolov5s only), v2.0, v3.0, v3.1, v4.0 and v5.0.
- For yolov5 v5.0, download .pt from yolov5 release v5.0,
git clone -b v5.0 https://github.com/ultralytics/yolov5.git
andgit clone -b yolov5-v5.0 https://github.com/wang-xinyu/tensorrtx.git
, then follow how-to-run in tensorrtx/yolov5-v5.0. - For yolov5 v4.0, download .pt from yolov5 release v4.0,
git clone -b v4.0 https://github.com/ultralytics/yolov5.git
andgit clone -b yolov5-v4.0 https://github.com/wang-xinyu/tensorrtx.git
, then follow how-to-run in tensorrtx/yolov5-v4.0. - For yolov5 v3.1, download .pt from yolov5 release v3.1,
git clone -b v3.1 https://github.com/ultralytics/yolov5.git
andgit clone -b yolov5-v3.1 https://github.com/wang-xinyu/tensorrtx.git
, then follow how-to-run in tensorrtx/yolov5-v3.1. - For yolov5 v3.0, download .pt from yolov5 release v3.0,
git clone -b v3.0 https://github.com/ultralytics/yolov5.git
andgit clone -b yolov5-v3.0 https://github.com/wang-xinyu/tensorrtx.git
, then follow how-to-run in tensorrtx/yolov5-v3.0. - For yolov5 v2.0, download .pt from yolov5 release v2.0,
git clone -b v2.0 https://github.com/ultralytics/yolov5.git
andgit clone -b yolov5-v2.0 https://github.com/wang-xinyu/tensorrtx.git
, then follow how-to-run in tensorrtx/yolov5-v2.0. - For yolov5 v1.0, download .pt from yolov5 release v1.0,
git clone -b v1.0 https://github.com/ultralytics/yolov5.git
andgit clone -b yolov5-v1.0 https://github.com/wang-xinyu/tensorrtx.git
, then follow how-to-run in tensorrtx/yolov5-v1.0.
- Choose the model s/m/l/x/s6/m6/l6/x6 from command line arguments.
- Input shape defined in yololayer.h
- Number of classes defined in yololayer.h, DO NOT FORGET TO ADAPT THIS, If using your own model
- INT8/FP16/FP32 can be selected by the macro in yolov5.cpp, INT8 need more steps, pls follow
How to Run
first and then go theINT8 Quantization
below - 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
- generate .wts from pytorch with .pt, or download .wts from model zoo
git clone -b v5.0 https://github.com/ultralytics/yolov5.git
git clone https://github.com/wang-xinyu/tensorrtx.git
// download https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt
cp {tensorrtx}/yolov5/gen_wts.py {ultralytics}/yolov5
cd {ultralytics}/yolov5
python gen_wts.py -w yolov5s.pt -o yolov5s.wts
// a file 'yolov5s.wts' will be generated.
- build tensorrtx/yolov5 and run
cd {tensorrtx}/yolov5/
// update CLASS_NUM in yololayer.h if your model is trained on custom dataset
mkdir build
cd build
cp {ultralytics}/yolov5/yolov5s.wts {tensorrtx}/yolov5/build
cmake ..
make
sudo ./yolov5 -s [.wts] [.engine] [s/m/l/x/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to plan file
sudo ./yolov5 -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed.
// For example yolov5s
sudo ./yolov5 -s yolov5s.wts yolov5s.engine s
sudo ./yolov5 -d yolov5s.engine ../samples
// For example Custom model with depth_multiple=0.17, width_multiple=0.25 in yolov5.yaml
sudo ./yolov5 -s yolov5_custom.wts yolov5.engine c 0.17 0.25
sudo ./yolov5 -d yolov5.engine ../samples
-
check the images generated, as follows. _zidane.jpg and _bus.jpg
-
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
-
Prepare calibration images, you can randomly select 1000s images from your train set. For coco, you can also download my calibration images
coco_calib
from GoogleDrive or BaiduPan pwd: a9wh -
unzip it in yolov5/build
-
set the macro
USE_INT8
in yolov5.cpp and make -
serialize the model and test
See the readme in home page.