blakeblackshear.frigate/docker/tensorrt_models.sh
Nate Meyer 3f05f74ecb
Nvidia TensorRT detector (#4718)
* Initial WIP dockerfile and scripts to add tensorrt support

* Add tensorRT detector

* WIP attempt to install TensorRT 8.5

* Updates to detector for cuda python library

* TensorRT Cuda library rework WIP

Does not run

* Fixes from rebase to detector factory

* Fix parsing output memory pointer

* Handle TensorRT logs with the python logger

* Use non-async interface and convert input data to float32. Detection runs without error.

* Make TensorRT a separate build from the base Frigate image.

* Add script and documentation for generating TRT Models

* Add support for TensorRT devcontainer

* Add labelmap to trt model script and docs.  Cleanup of old scripts.

* Update detect to normalize input tensor using model input type

* Add config for selecting GPU. Fix Async inference. Update documentation.

* Update some CUDA libraries to clean up version warning

* Add CI stage to build TensorRT tag

* Add note in docs for image tag and model support
2022-12-30 10:53:17 -06:00

37 lines
1.1 KiB
Bash
Executable File

#!/bin/bash
set -euxo pipefail
CUDA_HOME=/usr/local/cuda
LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64
OUTPUT_FOLDER=/tensorrt_models
echo "Generating the following TRT Models: ${YOLO_MODELS:="yolov4-tiny-288,yolov4-tiny-416,yolov7-tiny-416"}"
# Create output folder
mkdir -p ${OUTPUT_FOLDER}
# Install packages
pip install --upgrade pip && pip install onnx==1.9.0 protobuf==3.20.3
# Clone tensorrt_demos repo
git clone --depth 1 https://github.com/yeahme49/tensorrt_demos.git /tensorrt_demos
# Build libyolo
cd /tensorrt_demos/plugins && make all
cp libyolo_layer.so ${OUTPUT_FOLDER}/libyolo_layer.so
# Download yolo weights
cd /tensorrt_demos/yolo && ./download_yolo.sh
# Build trt engine
cd /tensorrt_demos/yolo
for model in ${YOLO_MODELS//,/ }
do
python3 yolo_to_onnx.py -m ${model}
python3 onnx_to_tensorrt.py -m ${model}
cp /tensorrt_demos/yolo/${model}.trt ${OUTPUT_FOLDER}/${model}.trt;
done
# Download Labelmap
wget -q https://github.com/openvinotoolkit/open_model_zoo/raw/master/data/dataset_classes/coco_91cl.txt -O ${OUTPUT_FOLDER}/coco_91cl.txt