tensorrt¶

def load(self) -> None:
    """Load TensorRT engine and prepare for inference.

    Performs:
        - Load engine from disk
        - Create execution context
        - Allocate device memory for inputs/outputs
        - Create CUDA stream
        - Warmup inference

    Raises:
        FileNotFoundError:  If engine file does not exist.
        RuntimeError: If TensorRT fails to deserialize engine.
    """
    logger.info(f"Loading TensorRT engine from {self.model_path}")

    # Load engine
    self.runtime = trt.Runtime(self.logger_trt)
    with self.model_path.open("rb") as f:
        engine_data = f.read()
    self.engine = self.runtime.deserialize_cuda_engine(engine_data)
    self.context = self.engine.create_execution_context()

    # Create CUDA stream
    self.stream = cuda.Stream()

    # Allocate memory for all bindings
    for i in range(self.engine.num_bindings):
        binding_shape = self.engine.get_binding_shape(i)
        dtype = trt.nptype(self.engine.get_binding_dtype(i))

        # Calculate size (reuse mixin)
        size = self._calculate_binding_size(binding_shape, dtype)

        # Allocate device memory
        mem = cuda.mem_alloc(size)
        self.bindings.append(int(mem))

        if self.engine.binding_is_input(i):
            self.inputs.append(mem)
        else:
            self.outputs.append(mem)

    logger.info(f"Engine loaded: inputs={len(self.inputs)}, outputs={len(self.outputs)}")

    # Warmup
    self._warmup()

def infer(self, input: np.ndarray) -> t.Any:
    """Run inference on a single input.

    Uses CUDA async operations for efficient processing.

    Args:
        input: Input numpy array.

    Returns:
        Output array with shape determined by model.

    Raises:
        RuntimeError: If engine is not loaded.

    Example:
        ```python
        with runtime:
            input = np.random.randn(1, 3, 640, 640).astype(
                np.float32
            )
            output = runtime.infer(input)
        ```
    """
    if self.context is None or self.stream is None:
        raise RuntimeError("Engine not loaded. Call load() first.")

    # Copy input to device
    cuda.memcpy_htod_async(self.inputs[0], input, self.stream)

    # Execute
    self.context.execute_async_v2(
        bindings=self.bindings,
        stream_handle=self.stream.handle,
    )

    # Copy output to host
    output_shape = self.engine.get_binding_shape(1)
    output_dtype = trt.nptype(self.engine.get_binding_dtype(1))
    output = np.empty(trt.volume(output_shape), dtype=output_dtype)

    cuda.memcpy_dtoh_async(output, self.outputs[0], self.stream)
    self.stream.synchronize()

    return output.reshape(output_shape)

def infer_batch(self, inputs: list[np.ndarray]) -> list[t.Any]:
    """Run inference on a batch of inputs.

    Concatenates inputs and runs batch inference, then splits
    the output back into individual results.

    Args:
        inputs: List of input arrays. Each should have shape (1, C, H, W).

    Returns:
        List of outputs, one per input. Each maintains batch dimension.

    Raises:
        RuntimeError: If engine is not loaded.

    Example:
        ```python
        with runtime:
            batch = [
                np.random.randn(1, 3, 640, 640).astype(np.float32),
                np.random.randn(1, 3, 640, 640).astype(np.float32),
            ]
            outputs = runtime.infer_batch(batch)
        ```
    """
    if self.context is None:
        raise RuntimeError("Engine not loaded. Call load() first.")

    # Concatenate and infer
    batch = np.concatenate(inputs, axis=0)
    batch_output = self.infer(batch)

    # Split outputs
    batch_size = len(inputs)
    if isinstance(batch_output, np.ndarray):
        return [batch_output[i : i + 1] for i in range(batch_size)]

    raise TypeError(f"Unexpected output type:  {type(batch_output)}")

def unload(self) -> None:
    """Unload engine and free resources.

    Performs:
        - Free CUDA device memory
        - Release engine and context

    Safe to call multiple times.
    """
    logger.info("Unloading engine")

    # Free CUDA memory
    for mem in self.inputs + self.outputs:
        mem.free()

    self.inputs.clear()
    self.outputs.clear()
    self.bindings.clear()

    self.context = None
    self.engine = None
    self.runtime = None
    self.stream = None

    logger.info("Engine unloaded")