Skip to content

backends

sleap_nn.inference.layers.backends

Runtime backends for inference layers.

Exported:

  • :class:ModelBackend — Protocol every backend implements.
  • :class:TorchBackend — PyTorch nn.Module runtime with optional compile / FP16 / Conv-BN fusion.
  • :class:ONNXBackend — ONNX Runtime backend. Wraps an exported .onnx file; peak finding is baked into the graph.
  • :class:TensorRTBackend — TensorRT backend (CUDA-only, requires tensorrt extra).

Modules:

Name Description
base

ModelBackend protocol — the contract every runtime backend implements.
onnx_backend

ONNXBackend — runs an exported ONNX model under the ModelBackend protocol.
tensorrt_backend

TensorRTBackend — runs a serialized TensorRT engine.
torch_backend

TorchBackend — wrap any nn.Module (or Lightning module) for inference.

Classes:

Name Description
ModelBackend

Runtime-agnostic forward-pass contract.
ONNXBackend

ONNX Runtime backend conforming to :class:ModelBackend.
TensorRTBackend

Native TensorRT engine backend conforming to :class:ModelBackend.
TorchBackend

PyTorch nn.Module backend with opt-in compile / FP16 / fusion.

ModelBackend

Bases: Protocol

Runtime-agnostic forward-pass contract.

Any object that satisfies this protocol can power any InferenceLayer subclass. Verify with isinstance(obj, ModelBackend) at construction time.

Methods:

Name Description
__call__

Run the model forward pass.
warmup

Run dummy forward passes to prime the backend.

Attributes:

Name Type Description
device str

Device the backend runs on. "cpu", "cuda", "cuda:0", "mps".
does_baked_postproc bool

True if this backend already performs peak finding internally.
Source code in sleap_nn/inference/layers/backends/base.py 
18
19
20
21
22
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
@runtime_checkable
class ModelBackend(Protocol):
    """Runtime-agnostic forward-pass contract.

    Any object that satisfies this protocol can power any
    ``InferenceLayer`` subclass. Verify with ``isinstance(obj, ModelBackend)``
    at construction time.
    """

    @property
    def device(self) -> str:
        """Device the backend runs on. ``"cpu"``, ``"cuda"``, ``"cuda:0"``, ``"mps"``."""
        ...

    @property
    def does_baked_postproc(self) -> bool:
        """``True`` if this backend already performs peak finding internally.

        ONNX and TensorRT export wrappers bake normalization + peak finding
        + (optionally) PAF scoring into the graph and return precomputed
        peaks. When this property is ``True``, the wrapping ``InferenceLayer``
        must skip its own Python-side peak finding and only apply coordinate
        transforms to whatever the backend returns.

        For pure PyTorch (``TorchBackend``) this is always ``False``.
        """
        ...

    def __call__(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Run the model forward pass.

        Args:
            x: Preprocessed input. Shape ``(B, C, H, W)``. The backend
                accepts whatever the upstream layer hands it — e.g.
                ``TorchBackend`` accepts uint8 because the Lightning
                module normalizes internally; an ONNX backend may also
                accept uint8 because normalization is baked.

        Returns:
            Dict of output tensors. Keys depend on the wrapped model:

            - Torch (``does_baked_postproc=False``)::

                {"SingleInstanceConfmapsHead": (B, N, H, W), ...}

            - ONNX/TRT (``does_baked_postproc=True``)::

                {"peaks": (B, I, N, 2), "peak_vals": (B, I, N), ...}
        """
        ...

    def warmup(self, input_shape: Tuple[int, ...]) -> None:
        """Run dummy forward passes to prime the backend.

        Particularly important on MPS (≈73× cold-start ratio per the
        benchmark suite) and CUDA-with-compile (where the first call
        triggers JIT compilation).

        Args:
            input_shape: Shape of the dummy tensor to allocate.
        """
        ...

device property

Device the backend runs on. "cpu", "cuda", "cuda:0", "mps".

does_baked_postproc property

True if this backend already performs peak finding internally.

ONNX and TensorRT export wrappers bake normalization + peak finding + (optionally) PAF scoring into the graph and return precomputed peaks. When this property is True, the wrapping InferenceLayer must skip its own Python-side peak finding and only apply coordinate transforms to whatever the backend returns.

For pure PyTorch (TorchBackend) this is always False.

__call__(x)

Run the model forward pass.

Parameters:

Name Type Description Default
x Tensor

Preprocessed input. Shape (B, C, H, W). The backend accepts whatever the upstream layer hands it — e.g. TorchBackend accepts uint8 because the Lightning module normalizes internally; an ONNX backend may also accept uint8 because normalization is baked.

 required

Returns:

Type Description
Dict[str, Tensor]

Dict of output tensors. Keys depend on the wrapped model:

  • Torch (does_baked_postproc=False)::

    {"SingleInstanceConfmapsHead": (B, N, H, W), ...}

  • ONNX/TRT (does_baked_postproc=True)::

    {"peaks": (B, I, N, 2), "peak_vals": (B, I, N), ...}
Source code in sleap_nn/inference/layers/backends/base.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def __call__(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Run the model forward pass.

    Args:
        x: Preprocessed input. Shape ``(B, C, H, W)``. The backend
            accepts whatever the upstream layer hands it — e.g.
            ``TorchBackend`` accepts uint8 because the Lightning
            module normalizes internally; an ONNX backend may also
            accept uint8 because normalization is baked.

    Returns:
        Dict of output tensors. Keys depend on the wrapped model:

        - Torch (``does_baked_postproc=False``)::

            {"SingleInstanceConfmapsHead": (B, N, H, W), ...}

        - ONNX/TRT (``does_baked_postproc=True``)::

            {"peaks": (B, I, N, 2), "peak_vals": (B, I, N), ...}
    """
    ...

warmup(input_shape)

Run dummy forward passes to prime the backend.

Particularly important on MPS (≈73× cold-start ratio per the benchmark suite) and CUDA-with-compile (where the first call triggers JIT compilation).

Parameters:

Name Type Description Default
input_shape Tuple[int, ...]

Shape of the dummy tensor to allocate.

 required
Source code in sleap_nn/inference/layers/backends/base.py 
69
70
71
72
73
74
75
76
77
78
79
def warmup(self, input_shape: Tuple[int, ...]) -> None:
    """Run dummy forward passes to prime the backend.

    Particularly important on MPS (≈73× cold-start ratio per the
    benchmark suite) and CUDA-with-compile (where the first call
    triggers JIT compilation).

    Args:
        input_shape: Shape of the dummy tensor to allocate.
    """
    ...

ONNXBackend

ONNX Runtime backend conforming to :class:ModelBackend.

Parameters:

Name Type Description Default
model_path

Path to an exported .onnx file.

 required
device

"cpu" / "cuda" / "auto" / "directml". Used to pick onnxruntime execution providers.

 required
providers

Explicit override for the execution-provider list. If None, providers are auto-selected from device.

 required
Notes

does_baked_postproc=True — the ONNX wrappers in sleap_nn/export/wrappers/ bake peak finding, normalization, and (top-down) crop extraction into the graph. Layer postprocess methods take the "peaks" / "peak_vals" keys directly from the session output instead of running Python peak finding.

Methods:

Name Description
__attrs_post_init__

Load the ONNX session and cache I/O metadata.
__call__

Run the ONNX session.
from_export_dir

Load an ONNX backend from an export directory containing model.onnx.
warmup

Run a single dummy forward to prime the runtime / GPU caches.

Attributes:

Name Type Description
does_baked_postproc bool

ONNX wrappers bake peak finding into the graph.
Source code in sleap_nn/inference/layers/backends/onnx_backend.py 
 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
166
167
168
169
170
171
172
173
174
175
176
177
@attrs.define(eq=False, slots=False)
class ONNXBackend:
    """ONNX Runtime backend conforming to :class:`ModelBackend`.

    Args:
        model_path: Path to an exported ``.onnx`` file.
        device: ``"cpu"`` / ``"cuda"`` / ``"auto"`` / ``"directml"``. Used
            to pick onnxruntime execution providers.
        providers: Explicit override for the execution-provider list. If
            ``None``, providers are auto-selected from ``device``.

    Notes:
        ``does_baked_postproc=True`` — the ONNX wrappers in
        ``sleap_nn/export/wrappers/`` bake peak finding, normalization,
        and (top-down) crop extraction into the graph. Layer postprocess
        methods take the ``"peaks"`` / ``"peak_vals"`` keys directly from
        the session output instead of running Python peak finding.
    """

    model_path: str
    device: str = "auto"
    providers: Optional[Iterable[str]] = None

    _session: object = attrs.field(default=None, init=False, repr=False)
    _input_name: str = attrs.field(default="", init=False, repr=False)
    _input_dtype: Optional[np.dtype] = attrs.field(default=None, init=False, repr=False)
    _output_names: list[str] = attrs.field(factory=list, init=False, repr=False)

    def __attrs_post_init__(self) -> None:
        """Load the ONNX session and cache I/O metadata."""
        try:
            import onnxruntime as ort
        except ImportError as exc:
            raise ImportError(
                "onnxruntime is required for ONNXBackend. Install with "
                "`pip install onnxruntime` (or `onnxruntime-gpu` for CUDA)."
            ) from exc

        if hasattr(ort, "preload_dlls"):
            # Auto-load CUDA/cuDNN libs from pip-installed nvidia-* packages.
            ort.preload_dlls()

        providers = (
            list(self.providers)
            if self.providers is not None
            else _select_providers(self.device, ort.get_available_providers())
        )
        self._session = ort.InferenceSession(self.model_path, providers=providers)

        in_info = self._session.get_inputs()[0]
        self._input_name = in_info.name
        self._input_dtype = _onnx_dtype_to_numpy(in_info.type)
        self._output_names = [out.name for out in self._session.get_outputs()]

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

    @property
    def does_baked_postproc(self) -> bool:
        """ONNX wrappers bake peak finding into the graph."""
        return True

    def __call__(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Run the ONNX session.

        Args:
            x: Input tensor. Cast to the session's expected dtype before
                handoff (most exported wrappers expect ``uint8``).

        Returns:
            Dict mapping the session's output names to torch tensors.
            Layer postprocess methods then look up ``"peaks"`` /
            ``"peak_vals"`` via the ``does_baked_postproc=True`` branch.
        """
        np_x = x.detach().cpu().numpy()
        if self._input_dtype is not None and np_x.dtype != self._input_dtype:
            np_x = np_x.astype(self._input_dtype)

        outputs = self._session.run(None, {self._input_name: np_x})
        return {
            name: torch.from_numpy(np.asarray(out))
            for name, out in zip(self._output_names, outputs)
        }

    def warmup(self, input_shape: Tuple[int, ...]) -> None:
        """Run a single dummy forward to prime the runtime / GPU caches."""
        dummy = np.zeros(input_shape, dtype=self._input_dtype or np.float32)
        self._session.run(None, {self._input_name: dummy})

    # ──────────────────────────────────────────────────────────────────
    # Convenience constructor
    # ──────────────────────────────────────────────────────────────────

    @classmethod
    def from_export_dir(
        cls, export_dir: Union[str, Path], device: str = "auto"
    ) -> "ONNXBackend":
        """Load an ONNX backend from an export directory containing ``model.onnx``.

        Args:
            export_dir: Directory written by ``sleap_nn export``.
            device: Device hint for execution-provider selection.

        Returns:
            A configured ``ONNXBackend``.
        """
        export_dir = Path(export_dir)
        candidates = sorted(export_dir.glob("*.onnx"))
        if not candidates:
            raise FileNotFoundError(f"No .onnx file found in {export_dir}")
        return cls(model_path=str(candidates[0]), device=device)

does_baked_postproc property

ONNX wrappers bake peak finding into the graph.

__attrs_post_init__()

Load the ONNX session and cache I/O metadata.

Source code in sleap_nn/inference/layers/backends/onnx_backend.py 
 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
def __attrs_post_init__(self) -> None:
    """Load the ONNX session and cache I/O metadata."""
    try:
        import onnxruntime as ort
    except ImportError as exc:
        raise ImportError(
            "onnxruntime is required for ONNXBackend. Install with "
            "`pip install onnxruntime` (or `onnxruntime-gpu` for CUDA)."
        ) from exc

    if hasattr(ort, "preload_dlls"):
        # Auto-load CUDA/cuDNN libs from pip-installed nvidia-* packages.
        ort.preload_dlls()

    providers = (
        list(self.providers)
        if self.providers is not None
        else _select_providers(self.device, ort.get_available_providers())
    )
    self._session = ort.InferenceSession(self.model_path, providers=providers)

    in_info = self._session.get_inputs()[0]
    self._input_name = in_info.name
    self._input_dtype = _onnx_dtype_to_numpy(in_info.type)
    self._output_names = [out.name for out in self._session.get_outputs()]

__call__(x)

Run the ONNX session.

Parameters:

Name Type Description Default
x Tensor

Input tensor. Cast to the session's expected dtype before handoff (most exported wrappers expect uint8).

 required

Returns:

Type Description
Dict[str, Tensor]

Dict mapping the session's output names to torch tensors. Layer postprocess methods then look up "peaks" / "peak_vals" via the does_baked_postproc=True branch.
Source code in sleap_nn/inference/layers/backends/onnx_backend.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def __call__(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Run the ONNX session.

    Args:
        x: Input tensor. Cast to the session's expected dtype before
            handoff (most exported wrappers expect ``uint8``).

    Returns:
        Dict mapping the session's output names to torch tensors.
        Layer postprocess methods then look up ``"peaks"`` /
        ``"peak_vals"`` via the ``does_baked_postproc=True`` branch.
    """
    np_x = x.detach().cpu().numpy()
    if self._input_dtype is not None and np_x.dtype != self._input_dtype:
        np_x = np_x.astype(self._input_dtype)

    outputs = self._session.run(None, {self._input_name: np_x})
    return {
        name: torch.from_numpy(np.asarray(out))
        for name, out in zip(self._output_names, outputs)
    }

from_export_dir(export_dir, device='auto') classmethod

Load an ONNX backend from an export directory containing model.onnx.

Parameters:

Name Type Description Default
export_dir Union[str, Path]

Directory written by sleap_nn export.

 required
device str

Device hint for execution-provider selection.

 'auto'

Returns:

Type Description
'ONNXBackend'

A configured ONNXBackend.
Source code in sleap_nn/inference/layers/backends/onnx_backend.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
@classmethod
def from_export_dir(
    cls, export_dir: Union[str, Path], device: str = "auto"
) -> "ONNXBackend":
    """Load an ONNX backend from an export directory containing ``model.onnx``.

    Args:
        export_dir: Directory written by ``sleap_nn export``.
        device: Device hint for execution-provider selection.

    Returns:
        A configured ``ONNXBackend``.
    """
    export_dir = Path(export_dir)
    candidates = sorted(export_dir.glob("*.onnx"))
    if not candidates:
        raise FileNotFoundError(f"No .onnx file found in {export_dir}")
    return cls(model_path=str(candidates[0]), device=device)

warmup(input_shape)

Run a single dummy forward to prime the runtime / GPU caches.

Source code in sleap_nn/inference/layers/backends/onnx_backend.py 
151
152
153
154
def warmup(self, input_shape: Tuple[int, ...]) -> None:
    """Run a single dummy forward to prime the runtime / GPU caches."""
    dummy = np.zeros(input_shape, dtype=self._input_dtype or np.float32)
    self._session.run(None, {self._input_name: dummy})

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)

TorchBackend

PyTorch nn.Module backend with opt-in compile / FP16 / fusion.

Parameters:

Name Type Description Default
model

The forward-pass owner. Typically a Lightning module (SingleInstanceLightningModule etc.) but any callable nn.Module works.

 required
device

"cpu", "cuda", "cuda:N", or "mps".

 required
use_compile

Wrap the model in torch.compile. CUDA-only. Emits a numeric-drift warning.

 required
compile_mode

Forwarded to torch.compile when enabled.

 required
use_fp16

Run the heavy forward ops in float16 via torch.autocast (CUDA only; fp32 master weights preserved, outputs cast back to fp32). Opt-in; benefits tensor-core hardware only. CUDA-only; counter-productive at batch < 4. Emits a drift warning.

 required
fuse_layers

Fold Conv2d → BatchNorm2d pairs in-place. Negligible speedup on the test UNets; opt-in.

 required
warmup_iterations

Number of dummy forwards to run inside :meth:warmup. 1 is enough on MPS / CUDA.

 required
Notes

slots=False is intentional — attrs-with-slots doesn't compose with Lightning's __getattr__ (which forwards to nn.Module). Using a regular class lets users mix the two without surprise.

Methods:

Name Description
__attrs_post_init__

Validate device / feature combination, fuse, and (optionally) compile.
__call__

Forward pass. Always returns a dict for protocol uniformity.
warmup

Prime the backend with warmup_iterations dummy forwards.

Attributes:

Name Type Description
does_baked_postproc bool

PyTorch returns raw confmaps; peak finding stays in Python.
Source code in sleap_nn/inference/layers/backends/torch_backend.py 
 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
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
@attrs.define(eq=False, slots=False)
class TorchBackend:
    """PyTorch ``nn.Module`` backend with opt-in compile / FP16 / fusion.

    Args:
        model: The forward-pass owner. Typically a Lightning module
            (``SingleInstanceLightningModule`` etc.) but any callable
            ``nn.Module`` works.
        device: ``"cpu"``, ``"cuda"``, ``"cuda:N"``, or ``"mps"``.
        use_compile: Wrap the model in ``torch.compile``. CUDA-only.
            Emits a numeric-drift warning.
        compile_mode: Forwarded to ``torch.compile`` when enabled.
        use_fp16: Run the heavy forward ops in float16 via ``torch.autocast``
            (CUDA only; fp32 master weights preserved, outputs cast back to
            fp32). Opt-in; benefits tensor-core hardware only.
            CUDA-only; counter-productive at batch < 4. Emits a drift warning.
        fuse_layers: Fold ``Conv2d → BatchNorm2d`` pairs in-place. Negligible
            speedup on the test UNets; opt-in.
        warmup_iterations: Number of dummy forwards to run inside
            :meth:`warmup`. ``1`` is enough on MPS / CUDA.

    Notes:
        ``slots=False`` is intentional — attrs-with-slots doesn't compose
        with Lightning's ``__getattr__`` (which forwards to ``nn.Module``).
        Using a regular class lets users mix the two without surprise.
    """

    model: nn.Module = attrs.field(repr=False)
    device: str = "cpu"
    use_compile: bool = False
    compile_mode: str = "reduce-overhead"
    use_fp16: bool = False
    fuse_layers: bool = False
    warmup_iterations: int = 1

    # Internal state, not part of the public surface.
    _compiled: Optional[nn.Module] = attrs.field(default=None, init=False, repr=False)

    def __attrs_post_init__(self) -> None:
        """Validate device / feature combination, fuse, and (optionally) compile."""
        self.model = self.model.to(self.device).eval()
        self._validate_device_features()

        if self.fuse_layers:
            self._fuse_conv_bn()

        # FP16 is applied at forward time via ``torch.autocast`` (CUDA only) —
        # see ``__call__``. We deliberately do NOT call ``model.half()``: a
        # destructive whole-model half cast raises an Input/weight dtype
        # mismatch on any ``forward`` that changes dtype internally (e.g.
        # ``image / 255.0`` normalization or an explicit ``.float()``) and gives
        # worse numerics than autocast's fp32 master weights. MPS/CPU keep fp32
        # (MPS has half kernels but no tensor cores — no win, warned above).

        if self.use_compile and self.device != "mps":
            self._compiled = torch.compile(
                self.model, mode=self.compile_mode, dynamic=False
            )

    # ──────────────────────────────────────────────────────────────────
    # Protocol surface
    # ──────────────────────────────────────────────────────────────────

    @property
    def does_baked_postproc(self) -> bool:
        """PyTorch returns raw confmaps; peak finding stays in Python."""
        return False

    def __call__(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Forward pass. Always returns a dict for protocol uniformity."""
        model = self._compiled if self._compiled is not None else self.model
        x = x.to(self.device, non_blocking=True)

        # FP16 (CUDA only) runs the heavy conv / matmul ops in half precision
        # via ``torch.autocast`` while keeping fp32 master weights and fp32
        # reductions. Autocast is robust to ``forward`` methods that change
        # dtype internally (``image / 255.0`` normalization, an explicit
        # ``.float()``, etc.) — a destructive ``model.half()`` would instead
        # raise an Input/weight dtype mismatch on those. Outputs are cast back
        # to fp32 below so downstream peak-finding sees fp32 confmaps.
        use_autocast = self.use_fp16 and "cuda" in self.device
        with torch.inference_mode():
            if use_autocast:
                with torch.autocast(device_type="cuda", dtype=torch.float16):
                    out = model(x)
            else:
                out = model(x)

        if isinstance(out, torch.Tensor):
            if use_autocast and out.dtype == torch.float16:
                out = out.float()
            return {"output": out}

        if isinstance(out, dict):
            if use_autocast:
                out = {
                    k: (
                        v.float()
                        if isinstance(v, torch.Tensor) and v.dtype == torch.float16
                        else v
                    )
                    for k, v in out.items()
                }
            return out

        raise TypeError(
            f"TorchBackend got unexpected output type {type(out).__name__}; "
            "expected Tensor or Dict[str, Tensor]."
        )

    def warmup(self, input_shape: Tuple[int, ...]) -> None:
        """Prime the backend with ``warmup_iterations`` dummy forwards."""
        if self.device == "cpu":
            return
        # Autocast handles the fp16 path internally, so the warmup dummy stays
        # fp32 (matching the dtype of the real preprocessed input).
        dummy = torch.zeros(input_shape, device=self.device, dtype=torch.float32)
        for _ in range(self.warmup_iterations):
            self(dummy)
        if "cuda" in self.device:
            torch.cuda.synchronize()
        elif self.device == "mps":
            torch.mps.synchronize()

    # ──────────────────────────────────────────────────────────────────
    # Internals
    # ──────────────────────────────────────────────────────────────────

    def _validate_device_features(self) -> None:
        """Warn-and-disable feature combinations that don't work cleanly.

        - MPS + ``torch.compile``: unreliable; force ``use_compile=False``.
        - MPS + FP16: kernels exist but tensor cores don't — no speedup;
          we keep ``use_fp16`` enabled but warn.
        - CUDA + ``torch.compile``: works; warn about graph fusion changing
          numerics.
        - CUDA + FP16: warn about ~4e-3 drift and small-batch regression.
        """
        if self.device == "mps":
            if self.use_compile:
                warnings.warn(
                    "torch.compile is unreliable on MPS; disabling.",
                    stacklevel=3,
                )
                self.use_compile = False
            if self.use_fp16:
                warnings.warn(
                    "FP16 on MPS gives no throughput gain (no tensor cores).",
                    stacklevel=3,
                )

        if "cuda" in self.device:
            if self.use_compile:
                warnings.warn(
                    "torch.compile changes numerics: graph fusion can substitute "
                    "TF32/reduced-precision kernels for FP32, producing small "
                    "drift vs. eager. Disable for parity-critical comparisons; "
                    "validate downstream metrics before shipping.",
                    stacklevel=3,
                )
            if self.use_fp16:
                warnings.warn(
                    "FP16 trades precision for speed. This backend runs the heavy "
                    "ops in half precision via torch.autocast (fp32 master weights "
                    "preserved). Measured max-abs-diff vs FP32 on the test "
                    "single-instance UNet (A40, batch 1-16): ~4e-3. Note: FP16 is "
                    "*counterproductive at small batch* — at batch=1 it ran 0.65× "
                    "the FP32 speed because tensor cores aren't saturated and "
                    "kernel-launch overhead dominates. Validate parity tests AND "
                    "speed at your actual batch size before enabling.",
                    stacklevel=3,
                )

    def _fuse_conv_bn(self) -> None:
        """Fold ``Conv2d → BatchNorm2d`` pairs in-place inside Sequentials.

        Fusion is restricted to ``nn.Sequential`` blocks, where execution
        order is guaranteed to match registration order.

        Why only ``nn.Sequential``? ``named_children()`` yields submodules in
        *registration* order, which is **not** the same as *execution* order
        for a module with a custom ``forward`` that reorders or skips
        submodules. Fusing a (Conv2d, BatchNorm2d) pair that happens to be
        registered consecutively — but is not actually applied consecutively
        in ``forward`` — folds the BN into the wrong conv and replaces the BN
        with ``nn.Identity()``, silently changing the model's output (observed
        max-abs-diff up to ~1.7 on a model whose ``forward`` reorders).

        ``nn.Sequential`` is the one container whose ``forward`` is *defined*
        to run children in registration order, so adjacency there genuinely
        implies "executed consecutively". We therefore only fuse Conv→BN pairs
        that are immediate neighbours inside an ``nn.Sequential`` and skip
        every other module — preferring a missed (~0% win) fusion over a
        silent mis-fuse. The sleap-nn UNet backbones build their Conv→BN
        stacks as ``nn.Sequential`` blocks, so the useful fusions are still
        covered.

        ``fuse_conv_bn_eval`` requires eval mode (running stats frozen); the
        backend has already called ``model.eval()`` before this runs.
        """
        from torch.nn.utils.fusion import fuse_conv_bn_eval

        def _fuse_in(parent: nn.Module) -> None:
            # Recurse into every child first so nested Sequentials get fused.
            for child in parent.children():
                _fuse_in(child)

            # Only Sequential guarantees registration order == execution
            # order, so it is the only place adjacency is safe to fuse.
            if not isinstance(parent, nn.Sequential):
                return

            entries: Any = list(parent.named_children())
            for i in range(len(entries) - 1):
                name, child = entries[i]
                bn_name, bn = entries[i + 1]
                if isinstance(child, nn.Conv2d) and isinstance(bn, nn.BatchNorm2d):
                    fused = fuse_conv_bn_eval(child, bn)
                    setattr(parent, name, fused)
                    setattr(parent, bn_name, nn.Identity())

        _fuse_in(self.model)

does_baked_postproc property

PyTorch returns raw confmaps; peak finding stays in Python.

__attrs_post_init__()

Validate device / feature combination, fuse, and (optionally) compile.

Source code in sleap_nn/inference/layers/backends/torch_backend.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def __attrs_post_init__(self) -> None:
    """Validate device / feature combination, fuse, and (optionally) compile."""
    self.model = self.model.to(self.device).eval()
    self._validate_device_features()

    if self.fuse_layers:
        self._fuse_conv_bn()

    # FP16 is applied at forward time via ``torch.autocast`` (CUDA only) —
    # see ``__call__``. We deliberately do NOT call ``model.half()``: a
    # destructive whole-model half cast raises an Input/weight dtype
    # mismatch on any ``forward`` that changes dtype internally (e.g.
    # ``image / 255.0`` normalization or an explicit ``.float()``) and gives
    # worse numerics than autocast's fp32 master weights. MPS/CPU keep fp32
    # (MPS has half kernels but no tensor cores — no win, warned above).

    if self.use_compile and self.device != "mps":
        self._compiled = torch.compile(
            self.model, mode=self.compile_mode, dynamic=False
        )

__call__(x)

Forward pass. Always returns a dict for protocol uniformity.

Source code in sleap_nn/inference/layers/backends/torch_backend.py 
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
def __call__(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Forward pass. Always returns a dict for protocol uniformity."""
    model = self._compiled if self._compiled is not None else self.model
    x = x.to(self.device, non_blocking=True)

    # FP16 (CUDA only) runs the heavy conv / matmul ops in half precision
    # via ``torch.autocast`` while keeping fp32 master weights and fp32
    # reductions. Autocast is robust to ``forward`` methods that change
    # dtype internally (``image / 255.0`` normalization, an explicit
    # ``.float()``, etc.) — a destructive ``model.half()`` would instead
    # raise an Input/weight dtype mismatch on those. Outputs are cast back
    # to fp32 below so downstream peak-finding sees fp32 confmaps.
    use_autocast = self.use_fp16 and "cuda" in self.device
    with torch.inference_mode():
        if use_autocast:
            with torch.autocast(device_type="cuda", dtype=torch.float16):
                out = model(x)
        else:
            out = model(x)

    if isinstance(out, torch.Tensor):
        if use_autocast and out.dtype == torch.float16:
            out = out.float()
        return {"output": out}

    if isinstance(out, dict):
        if use_autocast:
            out = {
                k: (
                    v.float()
                    if isinstance(v, torch.Tensor) and v.dtype == torch.float16
                    else v
                )
                for k, v in out.items()
            }
        return out

    raise TypeError(
        f"TorchBackend got unexpected output type {type(out).__name__}; "
        "expected Tensor or Dict[str, Tensor]."
    )

warmup(input_shape)

Prime the backend with warmup_iterations dummy forwards.

Source code in sleap_nn/inference/layers/backends/torch_backend.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
def warmup(self, input_shape: Tuple[int, ...]) -> None:
    """Prime the backend with ``warmup_iterations`` dummy forwards."""
    if self.device == "cpu":
        return
    # Autocast handles the fp16 path internally, so the warmup dummy stays
    # fp32 (matching the dtype of the real preprocessed input).
    dummy = torch.zeros(input_shape, device=self.device, dtype=torch.float32)
    for _ in range(self.warmup_iterations):
        self(dummy)
    if "cuda" in self.device:
        torch.cuda.synchronize()
    elif self.device == "mps":
        torch.mps.synchronize()