Skip to content

Latest commit

 

History

History

yolov5

README.md

yolov5

The Pytorch implementation is ultralytics/yolov5.

Different versions of yolov5

Currently, we support yolov5 v6.1.

  • For yolov5 v6.1, download .pt from yolov5 release v6.1, git clone -b v6.1 https://github.com/ultralytics/yolov5.git and git clone -b yolov5-v6.1 https://github.com/wang-xinyu/tensorrtx.git, then follow how-to-run in current page.

Config

  • Choose the model n/s/m/l/x/n6/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 the INT8 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
  • Optional, you can generate engine file with different batch/class_num/input_size by set command line parameter, without rebuild executable each time.

How to Run, yolov5s as example

  1. generate .wts from pytorch with .pt, or download .wts from model zoo
// clone code according to above #Different versions of yolov5
// download https://github.com/ultralytics/yolov5/releases/download/v6.1/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.
  1. 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] [n/s/m/l/x/n6/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

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

  2. 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
  1. optional, you can generate engine file with different batch/class_num/input_size by set command line parameter, without rebuild executable each time.

Just append -S b,c,h,w. Such as, to generate batch=2, classnum=18, input_h=384, input_w=640, just append -S 2,18,384,640.

INT8 Quantization

  1. 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

  2. unzip it in yolov5/build

  3. set the macro USE_INT8 in yolov5.cpp and make

  4. serialize the model and test

Reference:

  1. https://github.com/wang-xinyu/tensorrtx

  2. https://github.com/ultralytics/yolov5