edgefirst-validator 4.2.1__py3-none-any.whl

edgefirst/validator/runners/tensorrt.py

@@ -0,0 +1,321 @@
+from __future__ import annotations
+
+import json
+import time
+from typing import TYPE_CHECKING, Any, Tuple
+
+import numpy as np
+
+from edgefirst.validator.runners.core import Runner
+
+if TYPE_CHECKING:
+    from edgefirst.validator.evaluators import ModelParameters, TimerContext
+
+
+class TensorRTRunner(Runner):
+    """
+    Loads and runs TensorRT Engines (.engine, .trt).  These models
+    are intended to be deployed on a device with a dedicated GPU.
+    This implementation was taken from the following sources:
+    https://github.com/NVIDIA/TensorRT/blob/main/samples/python/efficientdet/infer.py
+    https://github.com/ultralytics/ultralytics/blob/main/ultralytics/nn/autobackend.py#L326
+
+    Parameters
+    ----------
+    model: Any
+        This is typically the path to the model or the loaded TensorRT model.
+    parameters: ModelParameters
+        These are the model parameters set from the command line.
+    metadata: dict
+        The model metadata which contains information for decoding
+        the model outputs.
+    timer: TimerContext
+        A timer object for handling validation timings for the model.
+
+    Raises
+    ------
+    ImportError
+        Missing tensorrt library.
+    FileNotFoundError
+        Raised if the path to the model does not exist.
+    """
+
+    def __init__(
+        self,
+        model: Any,
+        parameters: ModelParameters,
+        metadata: dict,
+        timer: TimerContext
+    ):
+        super(TensorRTRunner, self).__init__(model, parameters, timer=timer)
+
+        try:
+            import tensorrt as trt
+        except ImportError:
+            raise ImportError(
+                "tensorrt is needed to run TensorRT models.")
+        try:
+            import pycuda.driver as cuda  # type: ignore
+        except ImportError:
+            raise ImportError(
+                "pycuda is needed to perform memory allocations for TensorRT.")
+
+        # Use autoprimaryctx if available (pycuda >= 2021.1) to
+        # prevent issues with other modules that rely on the primary
+        # device context.
+        try:
+            import pycuda.autoprimaryctx  # type: ignore
+        except ModuleNotFoundError:
+            try:
+                import pycuda.autoinit  # type: ignore
+            except ImportError:
+                raise ImportError(
+                    "supported NVIDIA GPU device is needed.")
+
+        # TensorRT are intended to run on the GPU.
+        self.parameters.engine = "gpu"
+        self.logger = trt.Logger(trt.Logger.INFO)
+        trt.init_libnvinfer_plugins(self.logger, namespace="")
+
+        # Read file
+        if isinstance(model, str):
+            with open(model, "rb") as f, trt.Runtime(self.logger) as runtime:
+                assert runtime
+                try:
+                    meta_len = int.from_bytes(
+                        f.read(4), byteorder="little")  # read metadata length
+                    metadata = json.loads(
+                        f.read(meta_len).decode("utf-8"))  # read metadata
+                    dla = metadata.get("dla", None)
+                    if dla is not None:
+                        runtime.DLA_core = int(dla)
+                except UnicodeDecodeError:
+                    # engine file may lack embedded Ultralytics metadata
+                    f.seek(0)
+                self.model = runtime.deserialize_cuda_engine(
+                    f.read())  # read engine
+                assert self.model
+
+        self.context = self.model.create_execution_context()
+        assert self.context
+
+        self.output_names = []
+        self.output = []
+        self.input = []
+        self.allocations = []
+        self.scales = []  # Float values for dequantizing INT8 outputs.
+
+        num = range(self.model.num_bindings if hasattr(self.model, "num_bindings")
+                    else self.model.num_io_tensors)
+
+        for i in num:
+            name = self.model.get_tensor_name(i)
+            dtype = np.dtype(trt.nptype(self.model.get_tensor_dtype(name)))
+            shape = tuple(self.context.get_tensor_shape(name))
+            is_input = self.model.get_tensor_mode(
+                name) == trt.TensorIOMode.INPUT
+            if is_input:
+                if -1 in tuple(self.model.get_tensor_shape(name)):
+                    self.context.set_input_shape(
+                        name, tuple(self.model.get_tensor_profile_shape(name, 0)[1]))
+            else:
+                self.output_names.append(name)
+
+                dyn_range = None
+                if hasattr(self.model, "get_tensor_dynamic_range"):
+                    dyn_range = self.model.get_tensor_dynamic_range(i)
+                elif hasattr(self.model, "get_binding_dynamic_range"):
+                    dyn_range = self.model.get_binding_dynamic_range(i)
+
+                if dyn_range and dyn_range > 0:
+                    scale = float(dyn_range) / 127.0
+                    self.scales.append(scale)
+
+            shape = tuple(self.context.get_tensor_shape(name))
+            size = dtype.itemsize
+            for s in shape:
+                size *= s
+            allocation = cuda.mem_alloc(size)
+            host_allocation = None if is_input else np.zeros(shape, dtype)
+
+            binding = {
+                "index": i,
+                "name": name,
+                "dtype": dtype,
+                "shape": list(shape),
+                "allocation": allocation,
+                "host_allocation": host_allocation,
+            }
+
+            self.allocations.append(allocation)
+            if is_input:
+                self.input.append(binding)
+            else:
+                self.output.append(binding)
+
+        assert len(self.input) > 0
+        assert len(self.output) > 0
+        assert len(self.allocations) > 0
+
+        outputs = [np.zeros(out["shape"], np.dtype(out["dtype"]))
+                   for out in self.output]
+        self.init_decoder(metadata=metadata, outputs=outputs)
+
+        if self.parameters.warmup > 0:
+            self.warmup()
+
+    def infer(self, image: np.ndarray) -> list:
+        """
+        Executes inference on a batch of images.
+
+        Parameters
+        ----------
+        image: np.ndarray
+            The image input after being preprocessed.
+            Typically this is an RGB image array with the same
+            input shape as the model.
+
+        Returns
+        -------
+        list
+            Raw model outputs stored inside a list.
+
+        Raises
+        ------
+        ImportError
+            Raised if the pycuda library is not installed.
+        """
+        try:
+            import pycuda.driver as cuda  # type: ignore
+        except ImportError:
+            raise ImportError(
+                "pycuda driver is needed for TensorRT inference.")
+
+        start = time.perf_counter()
+        # Copy I/O and Execute.
+        image = np.ascontiguousarray(image)
+        cuda.memcpy_htod(self.input[0]['allocation'], image)
+        elapsed = time.perf_counter() - start
+        # Setting the tensor as part of the preprocess time.
+        self.timer.add_time("input", elapsed * 1e3)  # Convert to ms.
+
+        with self.timer.time("inference"):
+            self.context.execute_v2(self.allocations)
+
+    def run_single_instance(self, image: np.ndarray) -> Any:
+        """
+        Run TensorRT inference on a single image and record the timings.
+        Memory copying to and from the GPU device be performed here.
+
+        Parameters
+        ----------
+        image: np.ndarray
+            The input image after being preprocessed.
+            Typically this is an RGB image array.
+
+        Returns
+        -------
+        Any
+            This could either return detection outputs after NMS.
+                np.ndarray
+                    The prediction bounding boxes.. [[box1], [box2], ...].
+                np.ndarray
+                    The prediction labels.. [cl1, cl2, ...].
+                np.ndarray
+                    The prediction confidence scores.. [score, score, ...]
+                    normalized between 0 and 1.
+            This could also return segmentation masks.
+                np.ndarray
+        """
+        try:
+            import pycuda.driver as cuda  # type: ignore
+        except ImportError:
+            raise ImportError(
+                "pycuda driver is needed for TensorRT inference.")
+
+        # Inference
+        self.infer(image)
+
+        start = time.perf_counter()
+        for o in range(len(self.output)):
+            cuda.memcpy_dtoh(
+                self.output[o]['host_allocation'],
+                self.output[o]['allocation'])
+
+        # Dequantize INT8 outputs if engine provides dynamic ranges.
+        if len(self.scales):
+            outputs = []
+            for i, out in enumerate(self.output):
+                o = out['host_allocation']
+                # only handle integer types here
+                if o.dtype in [np.int8, np.uint8]:
+                    o = o.astype(np.float32) * self.scales[i]
+                outputs.append(o)
+        else:
+            outputs = [o['host_allocation'] for o in self.output]
+        elapsed = time.perf_counter() - start
+
+        # Postprocessing
+        outputs = self.postprocessing(outputs)
+
+        # Fetching the tensor as part of the postprocess time.
+        self.timer.add_time("output", elapsed * 1e3)  # Convert to ms.
+
+        return outputs
+
+    def input_spec(self) -> Tuple[tuple, np.dtype]:
+        """
+        Grabs the specs for the input tensor
+        of the network. Useful to prepare memory allocations.
+
+        Returns
+        -------
+        shape: tuple
+            The shape of the input tensor.
+        dtype: np.dtype
+            The input datatype.
+        """
+        return self.input[0]['shape'], self.input[0]['dtype']
+
+    def output_spec(self) -> list:
+        """
+        Grabs the specs for the output tensors of the network.
+        Useful to prepare memory allocations.
+
+        Returns
+        -------
+        list
+            A list with two items per element, the shape and (numpy)
+            datatype of each output tensor.
+        """
+        specs = list()
+        for o in self.output:
+            specs.append((o['shape'], o['dtype']))
+        return specs
+
+    def get_input_type(self) -> np.dtype:
+        """
+        This returns the input type of the model for the
+        input with shape in the form
+        (batch size, channels, height, width) or
+        (batch size, height, width, channels).
+
+        Returns
+        -------
+        np.dtype
+            The input type of the model.
+        """
+        return self.input[0]['dtype']
+
+    def get_input_shape(self) -> list:
+        """
+        This fetches the model input shape.
+
+        Returns
+        -------
+        list
+            The model input shape
+            [batch size, channels, height, width] or
+            [batch size, height, width, channels].
+        """
+        return self.input[0]['shape']

edgefirst/validator/runners/tflite.py

@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+import os
+import time
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from edgefirst.validator.publishers.utils.logger import logger
+from edgefirst.validator.runners.core import Runner
+
+if TYPE_CHECKING:
+    from edgefirst.validator.evaluators import ModelParameters, TimerContext
+
+
+class TFliteRunner(Runner):
+    """
+    Loads and runs TensorFlow Lite models for inference.
+
+    Parameters
+    ----------
+    model: Any
+        The is typically the path to the model or the loaded TFLite model.
+    parameters: ModelParameters
+        These are the model parameters set from the command line.
+    metadata: dict
+        The model metadata which contains information for decoding
+        the model outputs.
+    timer: TimerContext
+        A timer object for handling validation timings for the model.
+
+    Raises
+    ------
+    ImportError
+        Raised if tflite_runtime and TensorFlow is not intalled.
+    FileNotFoundError
+        Raised if the path to the model does not exist.
+    """
+
+    def __init__(
+        self,
+        model: Any,
+        parameters: ModelParameters,
+        metadata: dict,
+        timer: TimerContext
+    ):
+        super(TFliteRunner, self).__init__(model, parameters, timer=timer)
+
+        try:
+            from tflite_runtime.interpreter import Interpreter  # type: ignore
+        except ImportError:
+            logger("tflite_runtime is not installed. Falling back to TensorFlow.",
+                   code="WARNING")
+            try:
+                import tensorflow as tf  # type: ignore
+                Interpreter = tf.lite.Interpreter
+            except ImportError:
+                raise ImportError(
+                    "TensorFlow or tflite_runtime is needed to run TFLite models.")
+
+        if isinstance(model, str):
+            if not os.path.exists(model):
+                raise FileNotFoundError(
+                    "The model '{}' does not exist.".format(model))
+
+            ext_delegate = self.select_delegates()
+            logger(f"Engine: {self.parameters.engine}", code="INFO")
+
+            if ext_delegate:
+                self.model = Interpreter(
+                    model_path=model,
+                    experimental_delegates=[ext_delegate]
+                )
+            else:
+                self.model = Interpreter(model_path=model)
+
+        self.model.allocate_tensors()
+
+        self.input_details = self.model.get_input_details()
+        self.output_details = self.model.get_output_details()
+        self.parameters.common.input_quantization = self.input_details[0][
+            "quantization"]
+        self.init_decoder(metadata=metadata, outputs=self.output_details)
+
+        # The model NumPy view object.
+        self.parameters.common.input_tensor = self.model.tensor(
+            self.input_details[0]["index"])
+
+        if self.parameters.warmup > 0:
+            self.warmup()
+
+    def select_delegates(self) -> Any:
+        """
+        Specify the delegates to load based on
+        the type of engine specified.
+
+        Returns
+        -------
+        Any
+            This is either the loaded delegate object or None
+            if it doesn't exist.
+        """
+        try:  # https://coral.ai/docs/edgetpu/tflite-python/#update-existing-tf-lite-code-for-the-edge-tpu
+            from tflite_runtime.interpreter import load_delegate  # type: ignore
+        except ImportError:
+            try:
+                import tensorflow as tf  # type: ignore
+                load_delegate = tf.lite.experimental.load_delegate
+            except ImportError:
+                raise ImportError(
+                    "TensorFlow or tflite_runtime is needed to run TFLite models.")
+
+        ext_delegate = None
+        if (os.path.exists(self.parameters.engine) and
+                self.parameters.engine.endswith(".so")):
+            ext_delegate = load_delegate(self.parameters.engine, {})
+        elif self.parameters.engine.lower() == "npu":
+            if os.path.exists("/usr/lib/libvx_delegate.so"):
+                self.parameters.engine = "/usr/lib/libvx_delegate.so"
+                ext_delegate = load_delegate(self.parameters.engine, {})
+                logger("Using '/usr/lib/libvx_delegate.so' for NPU inference.",
+                       code="INFO")
+            elif os.path.exists("/usr/lib/libneutron_delegate.so"):
+                self.parameters.engine = "/usr/lib/libneutron_delegate.so"
+                ext_delegate = load_delegate(self.parameters.engine, {})
+                logger("Using '/usr/lib/libneutron_delegate.so' for NPU inference.",
+                       code="INFO")
+            else:
+                logger(
+                    "Specified NPU, but cannot find '/usr/lib/lib<>_delegate.so'. " +
+                    "Specify the path to libvx_delegate.so or libneutron_delegate.so " +
+                    "in your system. Falling back to use the CPU instead.", code="WARNING")
+                self.parameters.engine = "cpu"
+        elif self.parameters.engine.lower() == "gpu":
+            logger(
+                "Inference with the GPU is currently not supported for TFLite. " +
+                "Falling back to use the CPU instead.", code="WARNING")
+            self.parameters.engine = "cpu"
+        return ext_delegate
+
+    def run_single_instance(self, image: np.ndarray = None) -> Any:
+        """
+        Run TFLite inference on a single image and record the timings.
+
+        Parameters
+        ----------
+        image: np.ndarray
+            The input image after being preprocessed.
+            Typically this is an RGB image array. This is by default None.
+            Currently setting the image tensor using the logic below.
+
+            .. code-block:: python
+
+                input_tensor = self.model.tensor(self.input_details[0]["index"])
+                np.copyto(input_tensor(), image)
+
+        Returns
+        -------
+        Any
+            This could either return detection outputs after NMS.
+                np.ndarray
+                    The prediction bounding boxes.. [[box1], [box2], ...].
+                np.ndarray
+                    The prediction labels.. [cl1, cl2, ...].
+                np.ndarray
+                    The prediction confidence scores.. [score, score, ...]
+                    normalized between 0 and 1.
+            This could also return segmentation masks.
+                np.ndarray
+        """
+
+        # Inference
+        with self.timer.time("inference"):
+            self.model.invoke()
+
+        start = time.perf_counter()
+        outputs = [self.model.get_tensor(output["index"])
+                   for output in self.output_details]
+        elapsed = time.perf_counter() - start
+
+        # Postprocessing
+        outputs = self.postprocessing(outputs)
+
+        # Fetching the tensor as part of the postprocess time.
+        self.timer.add_time("output", elapsed * 1e3)  # Convert to ms.
+
+        return outputs
+
+    def get_input_type(self) -> np.dtype:
+        """
+        This returns the input type of the model with shape
+        (batch size, channels, height, width) or
+        (batch size, height, width, channels).
+
+        Returns
+        -------
+        np.dtype
+            The input type of the model.
+        """
+        return self.input_details[0]["dtype"]
+
+    def get_input_shape(self) -> np.ndarray:
+        """
+        Grabs the model input shape.
+
+        Returns
+        -------
+        np.ndarray
+            The model input shape (batch size, channels, height, width) or
+            (batch size, height, width, channels).
+        """
+        for input in self.input_details:
+            shape = input["shape"]
+            # Detection shapes are in the form (batch size, height, width,
+            # channels=3).
+            if len(shape) == 4:
+                if shape[1] == 3 or shape[-1] == [3]:
+                    return shape
+        # If it does not conform with expected format, return the first
+        # element.
+        return self.input_details[0]["shape"]