Mirrors/blakeblackshear.frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2025-07-26 13:47:03 +02:00
Code Issues Projects Releases Wiki Activity

Add support for TensorRT v10 (multiple api calls have changed) (#11166)

Browse Source 
* Add support for TensorRT v10 (multiple api calls have changed)

* Remove unnecessary size check in TensorRT v10 block

* Refactor to reduce code duplication

* Fix wrong function name in new _get_binding_dtype function and only return input check (not assertion) in new _binding_is_input function

* Add space around TRT_VERSION variable assignment (=) to respect linting

* More linting fix

* Update frigate/detectors/plugins/tensorrt.py

Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com>

* More linting

---------

Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com>
This commit is contained in:
Rémi Bédard-Couture 2024-05-22 08:57:52 -04:00 committed by GitHub
parent e91f3d8d9b
commit 592b645231
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 37 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

48
frigate/detectors/plugins/tensorrt.py
View File

@ -7,6 +7,8 @@ try:
import tensorrt as trt import tensorrt as trt
from cuda import cuda from cuda import cuda
TRT_VERSION = int(trt.__version__[0 : trt.__version__.find(".")])
TRT_SUPPORT = True TRT_SUPPORT = True
except ModuleNotFoundError: except ModuleNotFoundError:
TRT_SUPPORT = False TRT_SUPPORT = False
@ -88,20 +90,46 @@ class TensorRtDetector(DetectionApi):
with open(model_path, "rb") as f, trt.Runtime(self.trt_logger) as runtime: with open(model_path, "rb") as f, trt.Runtime(self.trt_logger) as runtime:
return runtime.deserialize_cuda_engine(f.read()) return runtime.deserialize_cuda_engine(f.read())
def _binding_is_input(self, binding):
if TRT_VERSION < 10:
return self.engine.binding_is_input(binding)
else:
return binding == "input"
def _get_binding_dims(self, binding):
if TRT_VERSION < 10:
return self.engine.get_binding_shape(binding)
else:
return self.engine.get_tensor_shape(binding)
def _get_binding_dtype(self, binding):
if TRT_VERSION < 10:
return self.engine.get_binding_dtype(binding)
else:
return self.engine.get_tensor_dtype(binding)
def _execute(self):
if TRT_VERSION < 10:
return self.context.execute_async_v2(
bindings=self.bindings, stream_handle=self.stream
)
else:
return self.context.execute_v2(self.bindings)
def _get_input_shape(self): def _get_input_shape(self):
"""Get input shape of the TensorRT YOLO engine.""" """Get input shape of the TensorRT YOLO engine."""
binding = self.engine[0] binding = self.engine[0]
assert self.engine.binding_is_input(binding) assert self._binding_is_input(binding)
binding_dims = self.engine.get_binding_shape(binding) binding_dims = self._get_binding_dims(binding)
if len(binding_dims) == 4: if len(binding_dims) == 4:
return ( return (
tuple(binding_dims[2:]), tuple(binding_dims[2:]),
trt.nptype(self.engine.get_binding_dtype(binding)), trt.nptype(self._get_binding_dtype(binding)),
) )
elif len(binding_dims) == 3: elif len(binding_dims) == 3:
return ( return (
tuple(binding_dims[1:]), tuple(binding_dims[1:]),
trt.nptype(self.engine.get_binding_dtype(binding)), trt.nptype(self._get_binding_dtype(binding)),
) )
else: else:
raise ValueError( raise ValueError(
@ -115,7 +143,7 @@ class TensorRtDetector(DetectionApi):
bindings = [] bindings = []
output_idx = 0 output_idx = 0
for binding in self.engine: for binding in self.engine:
binding_dims = self.engine.get_binding_shape(binding) binding_dims = self._get_binding_dims(binding)
if len(binding_dims) == 4: if len(binding_dims) == 4:
# explicit batch case (TensorRT 7+) # explicit batch case (TensorRT 7+)
size = trt.volume(binding_dims) size = trt.volume(binding_dims)
@ -126,21 +154,21 @@ class TensorRtDetector(DetectionApi):
raise ValueError( raise ValueError(
"bad dims of binding %s: %s" % (binding, str(binding_dims)) "bad dims of binding %s: %s" % (binding, str(binding_dims))
) )
nbytes = size * self.engine.get_binding_dtype(binding).itemsize nbytes = size * self._get_binding_dtype(binding).itemsize
# Allocate host and device buffers # Allocate host and device buffers
err, host_mem = cuda.cuMemHostAlloc( err, host_mem = cuda.cuMemHostAlloc(
nbytes, Flags=cuda.CU_MEMHOSTALLOC_DEVICEMAP nbytes, Flags=cuda.CU_MEMHOSTALLOC_DEVICEMAP
) )
assert err is cuda.CUresult.CUDA_SUCCESS, f"cuMemAllocHost returned {err}" assert err is cuda.CUresult.CUDA_SUCCESS, f"cuMemAllocHost returned {err}"
logger.debug( logger.debug(
f"Allocated Tensor Binding {binding} Memory {nbytes} Bytes ({size} * {self.engine.get_binding_dtype(binding)})" f"Allocated Tensor Binding {binding} Memory {nbytes} Bytes ({size} * {self._get_binding_dtype(binding)})"
) )
err, device_mem = cuda.cuMemAlloc(nbytes) err, device_mem = cuda.cuMemAlloc(nbytes)
assert err is cuda.CUresult.CUDA_SUCCESS, f"cuMemAlloc returned {err}" assert err is cuda.CUresult.CUDA_SUCCESS, f"cuMemAlloc returned {err}"
# Append the device buffer to device bindings. # Append the device buffer to device bindings.
bindings.append(int(device_mem)) bindings.append(int(device_mem))
# Append to the appropriate list. # Append to the appropriate list.
if self.engine.binding_is_input(binding): if self._binding_is_input(binding):
logger.debug(f"Input has Shape {binding_dims}") logger.debug(f"Input has Shape {binding_dims}")
inputs.append(HostDeviceMem(host_mem, device_mem, nbytes, size)) inputs.append(HostDeviceMem(host_mem, device_mem, nbytes, size))
else: else:
@ -170,9 +198,7 @@ class TensorRtDetector(DetectionApi):
] ]
# Run inference. # Run inference.
if not self.context.execute_async_v2( if not self._execute():
bindings=self.bindings, stream_handle=self.stream
):
logger.warn("Execute returned false") logger.warn("Execute returned false")
# Transfer predictions back from the GPU. # Transfer predictions back from the GPU.