Optimize OpenVINO and ONNX Model Runners (#20063)

* Use re-usable inference request to reduce CPU usage

* Share tensor

* Don't count performance

* Create openvino runner class

* Break apart onnx runner

* Add specific note about inability to use CUDA graphs for some models

* Adjust rknn to use RKNNRunner

* Use optimized runner

* Add support for non-complex models for CudaExecutionProvider

* Use core mask for rknn

* Correctly handle cuda input

* Cleanup

* Sort imports

frigate/detectors/plugins/onnx.py

@@ -6,6 +6,7 @@ from pydantic import Field
 from typing_extensions import Literal
 
 from frigate.detectors.detection_api import DetectionApi
+from frigate.detectors.detection_runners import CudaGraphRunner
 from frigate.detectors.detector_config import (
     BaseDetectorConfig,
     ModelTypeEnum,
@@ -23,53 +24,6 @@ logger = logging.getLogger(__name__)
 DETECTOR_KEY = "onnx"
 
 
-class CudaGraphRunner:
-    """Encapsulates CUDA Graph capture and replay using ONNX Runtime IOBinding.
-
-    This runner assumes a single tensor input and binds all model outputs.
-    """
-
-    def __init__(self, session: ort.InferenceSession, cuda_device_id: int):
-        self._session = session
-        self._cuda_device_id = cuda_device_id
-        self._captured = False
-        self._io_binding: ort.IOBinding | None = None
-        self._input_name: str | None = None
-        self._output_names: list[str] | None = None
-        self._input_ortvalue: ort.OrtValue | None = None
-        self._output_ortvalues: ort.OrtValue | None = None
-
-    def run(self, input_name: str, tensor_input: np.ndarray):
-        tensor_input = np.ascontiguousarray(tensor_input)
-
-        if not self._captured:
-            # Prepare IOBinding with CUDA buffers and let ORT allocate outputs on device
-            self._io_binding = self._session.io_binding()
-            self._input_name = input_name
-            self._output_names = [o.name for o in self._session.get_outputs()]
-
-            self._input_ortvalue = ort.OrtValue.ortvalue_from_numpy(
-                tensor_input, "cuda", self._cuda_device_id
-            )
-            self._io_binding.bind_ortvalue_input(self._input_name, self._input_ortvalue)
-
-            for name in self._output_names:
-                # Bind outputs to CUDA and allow ORT to allocate appropriately
-                self._io_binding.bind_output(name, "cuda", self._cuda_device_id)
-
-            # First IOBinding run to allocate, execute, and capture CUDA Graph
-            ro = ort.RunOptions()
-            self._session.run_with_iobinding(self._io_binding, ro)
-            self._captured = True
-            return self._io_binding.copy_outputs_to_cpu()
-
-        # Replay using updated input, copy results to CPU
-        self._input_ortvalue.update_inplace(tensor_input)
-        ro = ort.RunOptions()
-        self._session.run_with_iobinding(self._io_binding, ro)
-        return self._io_binding.copy_outputs_to_cpu()
-
-
 class ONNXDetectorConfig(BaseDetectorConfig):
     type: Literal[DETECTOR_KEY]
     device: str = Field(default="AUTO", title="Device Type")
@@ -114,7 +68,6 @@ class ONNXDetector(DetectionApi):
 
         try:
             if "CUDAExecutionProvider" in providers:
-                cuda_idx = providers.index("CUDAExecutionProvider")
                 self._cuda_device_id = options[cuda_idx].get("device_id", 0)
 
                 if options[cuda_idx].get("enable_cuda_graph"):
@@ -142,7 +95,7 @@ class ONNXDetector(DetectionApi):
         if self._cg_runner is not None:
             try:
                 # Run using CUDA graphs if available
-                tensor_output = self._cg_runner.run(model_input_name, tensor_input)
+                tensor_output = self._cg_runner.run({model_input_name: tensor_input})
             except Exception as e:
                 logger.warning(f"CUDA Graphs failed, falling back to regular run: {e}")
                 self._cg_runner = None

frigate/detectors/plugins/openvino.py

@@ -1,5 +1,4 @@
 import logging
-import os
 
 import numpy as np
 import openvino as ov
@@ -7,6 +6,7 @@ from pydantic import Field
 from typing_extensions import Literal
 
 from frigate.detectors.detection_api import DetectionApi
+from frigate.detectors.detection_runners import OpenVINOModelRunner
 from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum
 from frigate.util.model import (
     post_process_dfine,
@@ -37,20 +37,21 @@ class OvDetector(DetectionApi):
 
     def __init__(self, detector_config: OvDetectorConfig):
         super().__init__(detector_config)
-        self.ov_core = ov.Core()
         self.ov_model_type = detector_config.model.model_type
 
         self.h = detector_config.model.height
         self.w = detector_config.model.width
 
-        if not os.path.isfile(detector_config.model.path):
-            logger.error(f"OpenVino model file {detector_config.model.path} not found.")
-            raise FileNotFoundError
-
-        self.interpreter = self.ov_core.compile_model(
-            model=detector_config.model.path, device_name=detector_config.device
+        self.runner = OpenVINOModelRunner(
+            model_path=detector_config.model.path, device=detector_config.device
         )
 
+        # For dfine models, also pre-allocate target sizes tensor
+        if self.ov_model_type == ModelTypeEnum.dfine:
+            self.target_sizes_tensor = ov.Tensor(
+                np.array([[self.h, self.w]], dtype=np.int64)
+            )
+
         self.model_invalid = False
 
         if self.ov_model_type not in self.supported_models:
@@ -60,8 +61,8 @@ class OvDetector(DetectionApi):
             self.model_invalid = True
 
         if self.ov_model_type == ModelTypeEnum.ssd:
-            model_inputs = self.interpreter.inputs
-            model_outputs = self.interpreter.outputs
+            model_inputs = self.runner.compiled_model.inputs
+            model_outputs = self.runner.compiled_model.outputs
 
             if len(model_inputs) != 1:
                 logger.error(
@@ -80,8 +81,8 @@ class OvDetector(DetectionApi):
                 self.model_invalid = True
 
         if self.ov_model_type == ModelTypeEnum.yolonas:
-            model_inputs = self.interpreter.inputs
-            model_outputs = self.interpreter.outputs
+            model_inputs = self.runner.compiled_model.inputs
+            model_outputs = self.runner.compiled_model.outputs
 
             if len(model_inputs) != 1:
                 logger.error(
@@ -104,7 +105,9 @@ class OvDetector(DetectionApi):
             self.output_indexes = 0
             while True:
                 try:
-                    tensor_shape = self.interpreter.output(self.output_indexes).shape
+                    tensor_shape = self.runner.compiled_model.output(
+                        self.output_indexes
+                    ).shape
                     logger.info(
                         f"Model Output-{self.output_indexes} Shape: {tensor_shape}"
                     )
@@ -129,39 +132,32 @@ class OvDetector(DetectionApi):
         ]
 
     def detect_raw(self, tensor_input):
-        infer_request = self.interpreter.create_infer_request()
-        # TODO: see if we can use shared_memory=True
-        input_tensor = ov.Tensor(array=tensor_input)
+        if self.model_invalid:
+            return np.zeros((20, 6), np.float32)
 
         if self.ov_model_type == ModelTypeEnum.dfine:
-            infer_request.set_tensor("images", input_tensor)
-            target_sizes_tensor = ov.Tensor(
-                np.array([[self.h, self.w]], dtype=np.int64)
-            )
-            infer_request.set_tensor("orig_target_sizes", target_sizes_tensor)
-            infer_request.infer()
+            # Use named inputs for dfine models
+            inputs = {
+                "images": tensor_input,
+                "orig_target_sizes": np.array([[self.h, self.w]], dtype=np.int64),
+            }
+            outputs = self.runner.run_with_named_inputs(inputs)
             tensor_output = (
-                infer_request.get_output_tensor(0).data,
-                infer_request.get_output_tensor(1).data,
-                infer_request.get_output_tensor(2).data,
+                outputs["output0"],
+                outputs["output1"],
+                outputs["output2"],
             )
             return post_process_dfine(tensor_output, self.w, self.h)
 
-        infer_request.infer(input_tensor)
+        # Run inference using the runner
+        outputs = self.runner.run(tensor_input)
 
         detections = np.zeros((20, 6), np.float32)
 
-        if self.model_invalid:
-            return detections
-        elif self.ov_model_type == ModelTypeEnum.rfdetr:
-            return post_process_rfdetr(
-                [
-                    infer_request.get_output_tensor(0).data,
-                    infer_request.get_output_tensor(1).data,
-                ]
-            )
+        if self.ov_model_type == ModelTypeEnum.rfdetr:
+            return post_process_rfdetr(outputs)
         elif self.ov_model_type == ModelTypeEnum.ssd:
-            results = infer_request.get_output_tensor(0).data[0][0]
+            results = outputs[0][0][0]
 
             for i, (_, class_id, score, xmin, ymin, xmax, ymax) in enumerate(results):
                 if i == 20:
@@ -176,7 +172,7 @@ class OvDetector(DetectionApi):
                 ]
             return detections
         elif self.ov_model_type == ModelTypeEnum.yolonas:
-            predictions = infer_request.get_output_tensor(0).data
+            predictions = outputs[0]
 
             for i, prediction in enumerate(predictions):
                 if i == 20:
@@ -195,16 +191,10 @@ class OvDetector(DetectionApi):
                 ]
             return detections
         elif self.ov_model_type == ModelTypeEnum.yologeneric:
-            out_tensor = []
-
-            for item in infer_request.output_tensors:
-                out_tensor.append(item.data)
-
-            return post_process_yolo(out_tensor, self.w, self.h)
+            return post_process_yolo(outputs, self.w, self.h)
         elif self.ov_model_type == ModelTypeEnum.yolox:
-            out_tensor = infer_request.get_output_tensor()
             # [x, y, h, w, box_score, class_no_1, ..., class_no_80],
-            results = out_tensor.data
+            results = outputs[0]
             results[..., :2] = (results[..., :2] + self.grids) * self.expanded_strides
             results[..., 2:4] = np.exp(results[..., 2:4]) * self.expanded_strides
             image_pred = results[0, ...]

frigate/detectors/plugins/rknn.py

@@ -10,6 +10,7 @@ from pydantic import Field
 
 from frigate.const import MODEL_CACHE_DIR
 from frigate.detectors.detection_api import DetectionApi
+from frigate.detectors.detection_runners import RKNNModelRunner
 from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum
 from frigate.util.model import post_process_yolo
 from frigate.util.rknn_converter import auto_convert_model
@@ -61,18 +62,18 @@ class Rknn(DetectionApi):
                     "For more information, see: https://docs.deci.ai/super-gradients/latest/LICENSE.YOLONAS.html"
                 )
 
-        from rknnlite.api import RKNNLite
-
-        self.rknn = RKNNLite(verbose=False)
-        if self.rknn.load_rknn(model_props["path"]) != 0:
-            logger.error("Error initializing rknn model.")
-        if self.rknn.init_runtime(core_mask=core_mask) != 0:
-            logger.error(
-                "Error initializing rknn runtime. Do you run docker in privileged mode?"
-            )
+        self.runner = RKNNModelRunner(
+            model_path=model_props["path"],
+            model_type=config.model.model_type.value
+            if config.model.model_type
+            else None,
+            core_mask=core_mask,
+        )
 
     def __del__(self):
-        self.rknn.release()
+        if hasattr(self, "runner") and self.runner:
+            # The runner's __del__ method will handle cleanup
+            pass
 
     def get_soc(self):
         try:
@@ -305,9 +306,7 @@ class Rknn(DetectionApi):
             )
 
     def detect_raw(self, tensor_input):
-        output = self.rknn.inference(
-            [
-                tensor_input,
-            ]
-        )
+        # Prepare input for the runner
+        inputs = {"input": tensor_input}
+        output = self.runner.run(inputs)
         return self.post_process(output)