Skip to content

tensorrt_backend

sleap_nn.inference.layers.backends.tensorrt_backend

TensorRTBackend — runs a serialized TensorRT engine.

Like :class:ONNXBackend but the runtime is native TensorRT instead of ONNX Runtime. Same protocol surface: does_baked_postproc=True because the wrappers in sleap_nn/export/wrappers/ produce engines whose output names are "peaks", "peak_vals", etc.

CUDA-only. Constructing on a non-CUDA host raises a clear error; import tensorrt is deferred to __attrs_post_init__ so importing this module is cheap (and the import error doesn't break test collection on CPU-only hosts).

Classes:

Name Description
TensorRTBackend

Native TensorRT engine backend conforming to :class:ModelBackend.

TensorRTBackend

Native TensorRT engine backend conforming to :class:ModelBackend.

Parameters:

Name Type Description Default
engine_path

Path to a serialized TRT engine file (.trt).

 required
device

Must be "cuda" or "auto". Other values raise.

 required
Notes

Constructing this backend imports tensorrt lazily. On a host without CUDA / tensorrt installed, the constructor raises with a clear pointer at the right install extra ([tensorrt]).

Methods:

Name Description
__attrs_post_init__

Load the TRT engine + create an execution context.
__call__

Execute the TRT engine on x (must be on CUDA already).
from_export_dir

Load a TRT backend from an export directory containing *.trt.
warmup

Run a single dummy forward to prime the engine + GPU caches.

Attributes:

Name Type Description
does_baked_postproc bool

TRT engines exported from our wrappers bake peak finding.
Source code in sleap_nn/inference/layers/backends/tensorrt_backend.py 
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
@attrs.define(eq=False, slots=False)
class TensorRTBackend:
    """Native TensorRT engine backend conforming to :class:`ModelBackend`.

    Args:
        engine_path: Path to a serialized TRT engine file (``.trt``).
        device: Must be ``"cuda"`` or ``"auto"``. Other values raise.

    Notes:
        Constructing this backend imports ``tensorrt`` lazily. On a host
        without CUDA / ``tensorrt`` installed, the constructor raises with
        a clear pointer at the right install extra (``[tensorrt]``).
    """

    engine_path: str
    device: str = "cuda"

    _engine: object = attrs.field(default=None, init=False, repr=False)
    _context: object = attrs.field(default=None, init=False, repr=False)
    _input_names: list[str] = attrs.field(factory=list, init=False, repr=False)
    _output_names: list[str] = attrs.field(factory=list, init=False, repr=False)
    _trt: object = attrs.field(default=None, init=False, repr=False)

    def __attrs_post_init__(self) -> None:
        """Load the TRT engine + create an execution context."""
        if self.device not in ("cuda", "auto"):
            raise ValueError(
                f"TensorRTBackend only supports CUDA; got device={self.device!r}"
            )
        try:
            import tensorrt as trt
        except ImportError as exc:
            raise ImportError(
                "tensorrt is required for TensorRTBackend. Install with "
                "`pip install sleap-nn[tensorrt]` (Linux/Windows only)."
            ) from exc
        if not torch.cuda.is_available():
            raise RuntimeError(
                "TensorRTBackend requires a CUDA device; none is available."
            )

        self._trt = trt
        engine_path = Path(self.engine_path)
        if not engine_path.exists():
            raise FileNotFoundError(f"TensorRT engine not found: {engine_path}")

        logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, "rb") as f:
            self._engine = trt.Runtime(logger).deserialize_cuda_engine(f.read())
        if self._engine is None:
            raise RuntimeError(f"Failed to load TensorRT engine: {engine_path}")

        self._context = self._engine.create_execution_context()
        for i in range(self._engine.num_io_tensors):
            name = self._engine.get_tensor_name(i)
            if self._engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self._input_names.append(name)
            else:
                self._output_names.append(name)

    # ──────────────────────────────────────────────────────────────────
    # ModelBackend protocol surface
    # ──────────────────────────────────────────────────────────────────

    @property
    def does_baked_postproc(self) -> bool:
        """TRT engines exported from our wrappers bake peak finding."""
        return True

    def __call__(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Execute the TRT engine on ``x`` (must be on CUDA already).

        Args:
            x: ``(B, C, H, W)`` input tensor on CUDA. Auto-cast to the
                engine's expected dtype (``uint8`` for most exported
                wrappers).

        Returns:
            Dict mapping engine output names to torch tensors on CUDA.
        """
        trt = self._trt
        input_name = self._input_names[0]
        expected = self._engine.get_tensor_dtype(input_name)

        x = x.to("cuda", non_blocking=True)
        if expected == trt.DataType.UINT8:
            if x.dtype != torch.uint8:
                x = x.to(torch.uint8)
        elif x.dtype == torch.uint8:
            x = x.to(torch.float32)
        x = x.contiguous()

        self._context.set_input_shape(input_name, tuple(x.shape))

        bindings: Dict[str, int] = {input_name: x.data_ptr()}
        outputs: Dict[str, torch.Tensor] = {}
        for name in self._output_names:
            shape = tuple(self._context.get_tensor_shape(name))
            dtype = self._trt_dtype_to_torch(self._engine.get_tensor_dtype(name))
            outputs[name] = torch.empty(shape, dtype=dtype, device="cuda")
            bindings[name] = outputs[name].data_ptr()
        for name, ptr in bindings.items():
            self._context.set_tensor_address(name, ptr)

        stream = torch.cuda.current_stream().cuda_stream
        if not self._context.execute_async_v3(stream):
            raise RuntimeError("TensorRT inference failed")
        torch.cuda.current_stream().synchronize()
        return outputs

    def warmup(self, input_shape: Tuple[int, ...]) -> None:
        """Run a single dummy forward to prime the engine + GPU caches."""
        trt = self._trt
        input_name = self._input_names[0]
        expected = self._engine.get_tensor_dtype(input_name)
        dtype = torch.uint8 if expected == trt.DataType.UINT8 else torch.float32
        dummy = torch.zeros(input_shape, dtype=dtype, device="cuda")
        self(dummy)

    # ──────────────────────────────────────────────────────────────────
    # Helpers
    # ──────────────────────────────────────────────────────────────────

    def _trt_dtype_to_torch(self, trt_dtype) -> torch.dtype:
        """Map a TRT dtype to its torch counterpart (defaults to float32)."""
        trt = self._trt
        mapping = {
            trt.DataType.FLOAT: torch.float32,
            trt.DataType.HALF: torch.float16,
            trt.DataType.INT32: torch.int32,
            trt.DataType.INT8: torch.int8,
            trt.DataType.BOOL: torch.bool,
        }
        return mapping.get(trt_dtype, torch.float32)

    @classmethod
    def from_export_dir(cls, export_dir: Union[str, Path]) -> "TensorRTBackend":
        """Load a TRT backend from an export directory containing ``*.trt``."""
        export_dir = Path(export_dir)
        candidates = sorted(export_dir.glob("*.trt"))
        if not candidates:
            raise FileNotFoundError(f"No .trt file found in {export_dir}")
        return cls(engine_path=str(candidates[0]))

does_baked_postproc property

TRT engines exported from our wrappers bake peak finding.

__attrs_post_init__()

Load the TRT engine + create an execution context.

Source code in sleap_nn/inference/layers/backends/tensorrt_backend.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def __attrs_post_init__(self) -> None:
    """Load the TRT engine + create an execution context."""
    if self.device not in ("cuda", "auto"):
        raise ValueError(
            f"TensorRTBackend only supports CUDA; got device={self.device!r}"
        )
    try:
        import tensorrt as trt
    except ImportError as exc:
        raise ImportError(
            "tensorrt is required for TensorRTBackend. Install with "
            "`pip install sleap-nn[tensorrt]` (Linux/Windows only)."
        ) from exc
    if not torch.cuda.is_available():
        raise RuntimeError(
            "TensorRTBackend requires a CUDA device; none is available."
        )

    self._trt = trt
    engine_path = Path(self.engine_path)
    if not engine_path.exists():
        raise FileNotFoundError(f"TensorRT engine not found: {engine_path}")

    logger = trt.Logger(trt.Logger.WARNING)
    with open(engine_path, "rb") as f:
        self._engine = trt.Runtime(logger).deserialize_cuda_engine(f.read())
    if self._engine is None:
        raise RuntimeError(f"Failed to load TensorRT engine: {engine_path}")

    self._context = self._engine.create_execution_context()
    for i in range(self._engine.num_io_tensors):
        name = self._engine.get_tensor_name(i)
        if self._engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
            self._input_names.append(name)
        else:
            self._output_names.append(name)

__call__(x)

Execute the TRT engine on x (must be on CUDA already).

Parameters:

Name Type Description Default
x Tensor

(B, C, H, W) input tensor on CUDA. Auto-cast to the engine's expected dtype (uint8 for most exported wrappers).

 required

Returns:

Type Description
Dict[str, Tensor]

Dict mapping engine output names to torch tensors on CUDA.
Source code in sleap_nn/inference/layers/backends/tensorrt_backend.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def __call__(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Execute the TRT engine on ``x`` (must be on CUDA already).

    Args:
        x: ``(B, C, H, W)`` input tensor on CUDA. Auto-cast to the
            engine's expected dtype (``uint8`` for most exported
            wrappers).

    Returns:
        Dict mapping engine output names to torch tensors on CUDA.
    """
    trt = self._trt
    input_name = self._input_names[0]
    expected = self._engine.get_tensor_dtype(input_name)

    x = x.to("cuda", non_blocking=True)
    if expected == trt.DataType.UINT8:
        if x.dtype != torch.uint8:
            x = x.to(torch.uint8)
    elif x.dtype == torch.uint8:
        x = x.to(torch.float32)
    x = x.contiguous()

    self._context.set_input_shape(input_name, tuple(x.shape))

    bindings: Dict[str, int] = {input_name: x.data_ptr()}
    outputs: Dict[str, torch.Tensor] = {}
    for name in self._output_names:
        shape = tuple(self._context.get_tensor_shape(name))
        dtype = self._trt_dtype_to_torch(self._engine.get_tensor_dtype(name))
        outputs[name] = torch.empty(shape, dtype=dtype, device="cuda")
        bindings[name] = outputs[name].data_ptr()
    for name, ptr in bindings.items():
        self._context.set_tensor_address(name, ptr)

    stream = torch.cuda.current_stream().cuda_stream
    if not self._context.execute_async_v3(stream):
        raise RuntimeError("TensorRT inference failed")
    torch.cuda.current_stream().synchronize()
    return outputs

from_export_dir(export_dir) classmethod

Load a TRT backend from an export directory containing *.trt.

Source code in sleap_nn/inference/layers/backends/tensorrt_backend.py 
158
159
160
161
162
163
164
165
@classmethod
def from_export_dir(cls, export_dir: Union[str, Path]) -> "TensorRTBackend":
    """Load a TRT backend from an export directory containing ``*.trt``."""
    export_dir = Path(export_dir)
    candidates = sorted(export_dir.glob("*.trt"))
    if not candidates:
        raise FileNotFoundError(f"No .trt file found in {export_dir}")
    return cls(engine_path=str(candidates[0]))

warmup(input_shape)

Run a single dummy forward to prime the engine + GPU caches.

Source code in sleap_nn/inference/layers/backends/tensorrt_backend.py 
133
134
135
136
137
138
139
140
def warmup(self, input_shape: Tuple[int, ...]) -> None:
    """Run a single dummy forward to prime the engine + GPU caches."""
    trt = self._trt
    input_name = self._input_names[0]
    expected = self._engine.get_tensor_dtype(input_name)
    dtype = torch.uint8 if expected == trt.DataType.UINT8 else torch.float32
    dummy = torch.zeros(input_shape, dtype=dtype, device="cuda")
    self(dummy)