Skip to content

Commit

Permalink
Update library loading for tensorrt (blakeblackshear#5087)
Browse files Browse the repository at this point in the history
* Update library loading for tensorrt

* Add symlink to libnvrtc
  • Loading branch information
NateMeyer committed Jan 14, 2023
1 parent 19d17c8 commit 60b2315
Show file tree
Hide file tree
Showing 4 changed files with 10 additions and 8 deletions.
4 changes: 3 additions & 1 deletion Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,9 @@ COPY --from=rootfs / /
# Frigate w/ TensorRT Support as separate image
FROM frigate AS frigate-tensorrt
RUN --mount=type=bind,from=trt-wheels,source=/trt-wheels,target=/deps/trt-wheels \
pip3 install -U /deps/trt-wheels/*.whl
pip3 install -U /deps/trt-wheels/*.whl && \
ln -s libnvrtc.so.11.2 /usr/local/lib/python3.9/dist-packages/nvidia/cuda_nvrtc/lib/libnvrtc.so && \
ldconfig

# Dev Container w/ TRT
FROM devcontainer AS devcontainer-trt
Expand Down
5 changes: 5 additions & 0 deletions docker/rootfs/etc/ld.so.conf.d/cuda_tensorrt.conf
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
/usr/local/lib/python3.9/dist-packages/nvidia/cudnn/lib
/usr/local/lib/python3.9/dist-packages/nvidia/cuda_runtime/lib
/usr/local/lib/python3.9/dist-packages/nvidia/cublas/lib
/usr/local/lib/python3.9/dist-packages/nvidia/cuda_nvrtc/lib
/usr/local/lib/python3.9/dist-packages/tensorrt
6 changes: 0 additions & 6 deletions frigate/detectors/plugins/tensorrt.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,12 +75,6 @@ class TensorRtDetector(DetectionApi):

def _load_engine(self, model_path):
try:
ctypes.cdll.LoadLibrary(
"/usr/local/lib/python3.9/dist-packages/nvidia/cuda_runtime/lib/libcudart.so.11.0"
)
ctypes.cdll.LoadLibrary(
"/usr/local/lib/python3.9/dist-packages/tensorrt/libnvinfer.so.8"
)
trt.init_libnvinfer_plugins(self.trt_logger, "")

ctypes.cdll.LoadLibrary("/trt-models/libyolo_layer.so")
Expand Down
3 changes: 2 additions & 1 deletion requirements-tensorrt.txt
Original file line number Diff line number Diff line change
Expand Up @@ -5,4 +5,5 @@ cuda-python == 11.7; platform_machine == 'x86_64'
cython == 0.29.*; platform_machine == 'x86_64'
nvidia-cuda-runtime-cu11 == 11.7.*; platform_machine == 'x86_64'
nvidia-cublas-cu11 == 11.11.*; platform_machine == 'x86_64'
nvidia-cudnn-cu11 == 8.7.*; platform_machine == 'x86_64'
nvidia-cudnn-cu11 == 8.7.*; platform_machine == 'x86_64'
nvidia-cuda-nvrtc-cu11 == 11.7.*; platform_machine == 'x86_64'

0 comments on commit 60b2315

Please sign in to comment.