https://github.com/piotrostr/infer-trt
Interface for TensorRT engines inference along with an example of YOLOv4 engine being used.
https://github.com/piotrostr/infer-trt
deep-learning inference-optimization object-detection tensorrt
Last synced: about 1 year ago
JSON representation
Interface for TensorRT engines inference along with an example of YOLOv4 engine being used.
- Host: GitHub
- URL: https://github.com/piotrostr/infer-trt
- Owner: piotrostr
- License: mit
- Created: 2021-03-23T14:29:30.000Z (almost 5 years ago)
- Default Branch: master
- Last Pushed: 2022-05-07T13:42:27.000Z (almost 4 years ago)
- Last Synced: 2025-01-01T08:32:15.252Z (about 1 year ago)
- Topics: deep-learning, inference-optimization, object-detection, tensorrt
- Language: Python
- Homepage:
- Size: 17.6 KB
- Stars: 2
- Watchers: 2
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# infer-trt
Interface for TensorRT engines inference along with an example of YOLOv4 engine being used.
### System Requirements
- amd64/linux architecture
- Nvidia GPU with Tensor cores
### Setup
To start one can fork/clone this repo containing the code.
```bash
git clone https://github.com/piotrostr/infer-trt
```
TensorRT setup is quite complex so it is adviced to use the pre-build nvcr.io container.
```bash
docker pull nvcr.io/nvidia/tensorrt:22.03-py3
```
The example uses the opencv library, which can be built using the official
instructions which can be found in the
official docs.
Building it even with 8-core-cpu takes quite long so I'd advise to install the binary from conda-forge.
```
conda install -y -c conda-forge/label/gcc7 opencv
```
### Usage
Normally, the `yolo.onnx` weights will probably be a result of training, if using PyTorch or TF2 it is easy to make an ONNX export. For the sake of this example, pre-trained weights can be curled or wget'd from the ONNX repo.
```
curl -o yolo.onnx https://github.com/onnx/models/blob/main/vision/object_detection_segmentation/yolov3/model/yolov3-10.onnx
```
In order to obtain the TensorRT engine for a given model the `trtexec` tool can
be used to make an export from onnx weights file.
```bash
./trtexec \
--onnx=./yolo.onnx \
--best \
--workspace=1024 \
--saveEngine=./yolo.trt \
--optShapes=input:1x3x416x416
```
The `trt_model.py` contains a base class to be used for inheritance. Once the
preprocesing and postprocessing methods are overriden to match the steps
required per given model, it is ready for inference.
with its high-level api:
```python
import cv2
import pycuda.autoinit
import tensorrt as trt
from yolo import YOLO
yolo = YOLO(trt.Logger())
img = cv2.imread('some_img.png')
labels, confidences, bboxes = yolo(img)
```
### License
MIT©