yolov5-tensorrt

June 1, 2021 · View on GitHub

port pytorch/onnx yolov5 model to run on a Jetson Nano

ipynb is for testing pytorch code and exporting onnx models using Google Colab

python code runs numpy/tensorrt implementation on Jetson Nano

├── python
│   ├── lib
|       ├── demo.py
|       ├── Processor.py
|       ├── Visualizer.py
|       ├── classes.py
|       └── models
|           ├── yolov5s-simple-32.trt
|           ├── yolov5s-simple-16.trt
|           └── yolov5s-simple.onnx
│   └── export_tensorrt.py
  • convert yolov5 onnx model to tensorrt
  • pre-process image
  • run inference against input using tensorrt engine
  • post process output (forward pass)
  • apply nms thresholding on candidate boxes
  • visualize results

compile onnx model to trt

python3 export_tensorrt.py --help
usage: export_tensorrt.py [-h] [-m MODEL] [-fp FLOATINGPOINT] [-o OUTPUT]

compile Onnx model to TensorRT

optional arguments:
  -h, --help            show this help message and exit
  -m MODEL, --model MODEL
                        onnx file location inside ./lib/models
  -fp FLOATINGPOINT, --floatingpoint FLOATINGPOINT
                        floating point precision. 16 or 32
  -o OUTPUT, --output OUTPUT
                        name of trt output file

run demo

python3 demo.py --image=./path/to/image.jpg --model=./path/to/model.trt

performance

for now, only testing initial inference performance
nms, and post processing are slow rn
modelfp precisioninput sizetime (ms)
small-simple32640x640x3221 ms
small-simple16640x640x3?

object probability