Reference for ultralytics/nn/backends/tensorrt.py#

class TensorRTBackend(BaseBackend):
    """NVIDIA TensorRT inference backend for GPU-accelerated deployment.

    Loads and runs inference with NVIDIA TensorRT serialized engines (.engine files). Supports both TensorRT 7-9 and
    TensorRT 10/11 APIs, dynamic input shapes, FP16 precision, and DLA core offloading.
    """

def forward(self, im: torch.Tensor) -> list[torch.Tensor]:
    """Run NVIDIA TensorRT inference with dynamic shape handling.

    Args:
        im (torch.Tensor): Input image tensor in BCHW format on the CUDA device.

    Returns:
        (list[torch.Tensor]): Model predictions as a list of tensors on the CUDA device.
    """
    if self.dynamic and im.shape != self.bindings["images"].shape:
        if self.is_trt10:
            ok = self.context.set_input_shape("images", im.shape)
        else:
            ok = self.context.set_binding_shape(self.model.get_binding_index("images"), im.shape)
        if not ok:  # the profile refused the shape, so the bindings below would describe the wrong engine state
            raise ValueError(f"input size {tuple(im.shape)} is outside the TensorRT optimization profile")
        self.bindings["images"] = self.bindings["images"]._replace(shape=im.shape)
        for name in self.output_names:
            shape = (
                self.context.get_tensor_shape(name)
                if self.is_trt10
                else self.context.get_binding_shape(self.model.get_binding_index(name))
            )
            self.bindings[name].data.resize_(tuple(shape))

    s = self.bindings["images"].shape
    assert im.shape == s, f"input size {im.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}"

    if self.graph is None:
        self.bindings["images"] = self.bindings["images"]._replace(data=im)
        if not self.context.execute_v2([binding.data.data_ptr() for binding in self.bindings.values()]):
            raise RuntimeError("TensorRT inference execution failed")
    else:
        self.bindings["images"].data.copy_(im)  # the capture reads this address, so the input must land in it
        self.graph.replay()
    return [self.bindings[x].data for x in sorted(self.output_names)]

def load_model(self, weight: str | Path) -> None:
    """Load an NVIDIA TensorRT engine from a serialized .engine file.

    Args:
        weight (str | Path): Path to the .engine file with optional embedded metadata.
    """
    LOGGER.info(f"Loading {weight} for TensorRT inference...")

    if IS_JETSON and check_version(PYTHON_VERSION, "<=3.8.10"):
        check_requirements("numpy==1.23.5")

    try:
        import tensorrt as trt
    except ImportError:
        check_tensorrt()
        import tensorrt as trt

    check_version(trt.__version__, ">=7.0.0", hard=True)
    check_version(trt.__version__, "!=10.2.0", msg="https://github.com/ultralytics/ultralytics/pull/24367")

    if self.device.type == "cpu":
        self.device = torch.device("cuda:0")

    from ultralytics.utils.export.engine import get_tensorrt_logger

    Binding = namedtuple("Binding", ("name", "dtype", "shape", "data"))
    logger = get_tensorrt_logger()

    # Read engine file
    offset, metadata = self.engine_header(weight)
    with open(weight, "rb") as f, torch.cuda.device(self.device):
        runtime = trt.Runtime(logger)
        f.seek(offset)  # skip the metadata header, if any, that precedes the engine
        if (dla := metadata.get("dla")) is not None:
            runtime.DLA_core = int(dla)
        engine = runtime.deserialize_cuda_engine(f.read())
        self.apply_metadata(metadata)
        try:
            self.context = engine.create_execution_context()  # TensorRT binds this to the current device
        except Exception:
            LOGGER.error("TensorRT model exported with a different version than expected\n")
            raise

    # Setup bindings
    self.bindings = OrderedDict()
    self.output_names = []
    self.fp16 = False
    self.dynamic = False
    # TensorRT 10 and 11 both drop the legacy binding API in favor of named I/O tensors
    self.is_trt10 = not hasattr(engine, "num_bindings")
    num = range(engine.num_io_tensors) if self.is_trt10 else range(engine.num_bindings)

    for i in num:
        if self.is_trt10:
            name = engine.get_tensor_name(i)
            dtype = trt.nptype(engine.get_tensor_dtype(name))
            is_input = engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT
            shape = tuple(engine.get_tensor_shape(name))
            profile_shape = tuple(engine.get_tensor_profile_shape(name, 0)[2]) if is_input else None
        else:
            name = engine.get_binding_name(i)
            dtype = trt.nptype(engine.get_binding_dtype(i))
            is_input = engine.binding_is_input(i)
            shape = tuple(engine.get_binding_shape(i))
            profile_shape = tuple(engine.get_profile_shape(0, i)[1]) if is_input else None

        if is_input:
            if -1 in shape:
                self.dynamic = True
                if self.is_trt10:
                    self.context.set_input_shape(name, profile_shape)
                else:
                    self.context.set_binding_shape(i, profile_shape)
            if dtype == np.float16:
                self.fp16 = True
        else:
            self.output_names.append(name)

        shape = (
            tuple(self.context.get_tensor_shape(name))
            if self.is_trt10
            else tuple(self.context.get_binding_shape(i))
        )
        im = torch.from_numpy(np.empty(shape, dtype=dtype)).to(self.device)
        self.bindings[name] = Binding(name, dtype, shape, im)

    # Replay the engine from one captured CUDA graph instead of relaunching every kernel, worth ~50 us per call.
    # A dynamic engine would recapture on every shape change, and DLA or an embedded NMS does host work mid-stream
    # that cannot be captured, so both keep `execute_v2`.
    host = dla is not None or metadata.get("args", {}).get("nms", False)
    self.graph = None
    if TORCH_1_10 and self.is_trt10 and not self.dynamic and not host:
        for name, binding in self.bindings.items():
            self.context.set_tensor_address(name, binding.data.data_ptr())
        stream, graph = torch.cuda.Stream(self.device), torch.cuda.CUDAGraph()
        with torch.cuda.stream(stream):  # selects the engine's device as well, which the capture records on
            ok = self.context.execute_async_v3(stream.cuda_stream)  # TensorRT allocates on its first run
            stream.synchronize()
            with torch.cuda.graph(graph, stream=stream):
                ok &= self.context.execute_async_v3(stream.cuda_stream)
        self.graph = graph if ok else None  # a refused enqueue records nothing and would replay stale outputs

    self.model = engine