Skip to content

export

sleap_nn.export

Export utilities for sleap-nn.

Modules:

Name Description
cli

CLI entry points for export workflows.
exporters

Exporters for serialized model formats.
metadata

Metadata helpers for exported models.
predictors

Predictors for exported models.
utils

Utilities for export workflows.
wrappers

ONNX/TensorRT export wrappers.

Classes:

Name Description
ExportMetadata

Metadata embedded or saved alongside exported models.
ONNXPredictor

ONNX Runtime inference with provider selection.
TensorRTPredictor

TensorRT inference for exported models.

Functions:

Name Description
build_bottomup_candidate_template

Build candidate template matching ONNX wrapper's line_scores ordering.
export_model

Export a model to the requested format.
export_to_onnx

Export a PyTorch model to ONNX.
export_to_tensorrt

Export a PyTorch model to TensorRT format.
load_exported_model

Load an exported model and return a predictor instance.

ExportMetadata dataclass

Metadata embedded or saved alongside exported models.

Methods:

Name Description
default_timestamp

Return an ISO timestamp for export.
from_dict

Load from dict.
load

Load from JSON file.
save

Save to JSON file.
to_dict

Convert to JSON-serializable dict.
Source code in sleap_nn/export/metadata.py 
 15
 16
 17
 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
 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
@dataclass
class ExportMetadata:
    """Metadata embedded or saved alongside exported models."""

    # Version info
    sleap_nn_version: str
    export_timestamp: str
    export_format: str  # "onnx" or "tensorrt"

    # Model info
    model_type: str  # "centroid", "centered_instance", "topdown", "bottomup"
    model_name: str
    checkpoint_path: str

    # Architecture
    backbone: str
    n_nodes: int
    n_edges: int
    node_names: List[str]
    edge_inds: List[Tuple[int, int]]

    # Input/output spec
    input_scale: float
    input_channels: int
    output_stride: int
    crop_size: Optional[Tuple[int, int]] = None

    # Export parameters
    max_instances: Optional[int] = None
    max_peaks_per_node: Optional[int] = None
    max_batch_size: int = 1
    precision: str = "fp32"

    # Preprocessing - input is uint8 [0,255], normalized internally to float32 [0,1]
    input_dtype: str = "uint8"
    normalization: str = "0_to_1"

    # Multiclass model fields (optional)
    n_classes: Optional[int] = None
    class_names: Optional[List[str]] = None

    # Training config reference
    training_config_embedded: bool = False
    training_config_hash: str = ""

    def to_dict(self) -> Dict[str, Any]:
        """Convert to JSON-serializable dict."""
        data = asdict(self)
        data["edge_inds"] = [list(pair) for pair in self.edge_inds]
        if self.crop_size is not None:
            data["crop_size"] = list(self.crop_size)
        return data

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ExportMetadata":
        """Load from dict."""
        edge_inds = [tuple(pair) for pair in data.get("edge_inds", [])]
        crop_size = data.get("crop_size")
        if crop_size is not None:
            crop_size = tuple(crop_size)
        return cls(
            sleap_nn_version=data.get("sleap_nn_version", ""),
            export_timestamp=data.get("export_timestamp", ""),
            export_format=data.get("export_format", ""),
            model_type=data.get("model_type", ""),
            model_name=data.get("model_name", ""),
            checkpoint_path=data.get("checkpoint_path", ""),
            backbone=data.get("backbone", ""),
            n_nodes=int(data.get("n_nodes", 0)),
            n_edges=int(data.get("n_edges", 0)),
            node_names=list(data.get("node_names", [])),
            edge_inds=edge_inds,
            input_scale=float(data.get("input_scale", 1.0)),
            input_channels=int(data.get("input_channels", 1)),
            output_stride=int(data.get("output_stride", 1)),
            crop_size=crop_size,
            max_instances=data.get("max_instances"),
            max_peaks_per_node=data.get("max_peaks_per_node"),
            max_batch_size=int(data.get("max_batch_size", 1)),
            precision=data.get("precision", "fp32"),
            input_dtype=data.get("input_dtype", "uint8"),
            normalization=data.get("normalization", "0_to_1"),
            n_classes=data.get("n_classes"),
            class_names=data.get("class_names"),
            training_config_embedded=bool(data.get("training_config_embedded", False)),
            training_config_hash=data.get("training_config_hash", ""),
        )

    def save(self, path: str | Path) -> None:
        """Save to JSON file."""
        path = Path(path)
        path.write_text(json.dumps(self.to_dict(), indent=2, sort_keys=True))

    @classmethod
    def load(cls, path: str | Path) -> "ExportMetadata":
        """Load from JSON file."""
        path = Path(path)
        data = json.loads(path.read_text())
        return cls.from_dict(data)

    @classmethod
    def default_timestamp(cls) -> str:
        """Return an ISO timestamp for export."""
        return datetime.now().isoformat()

default_timestamp() classmethod

Return an ISO timestamp for export.

Source code in sleap_nn/export/metadata.py 
115
116
117
118
@classmethod
def default_timestamp(cls) -> str:
    """Return an ISO timestamp for export."""
    return datetime.now().isoformat()

from_dict(data) classmethod

Load from dict.

Source code in sleap_nn/export/metadata.py 
 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
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "ExportMetadata":
    """Load from dict."""
    edge_inds = [tuple(pair) for pair in data.get("edge_inds", [])]
    crop_size = data.get("crop_size")
    if crop_size is not None:
        crop_size = tuple(crop_size)
    return cls(
        sleap_nn_version=data.get("sleap_nn_version", ""),
        export_timestamp=data.get("export_timestamp", ""),
        export_format=data.get("export_format", ""),
        model_type=data.get("model_type", ""),
        model_name=data.get("model_name", ""),
        checkpoint_path=data.get("checkpoint_path", ""),
        backbone=data.get("backbone", ""),
        n_nodes=int(data.get("n_nodes", 0)),
        n_edges=int(data.get("n_edges", 0)),
        node_names=list(data.get("node_names", [])),
        edge_inds=edge_inds,
        input_scale=float(data.get("input_scale", 1.0)),
        input_channels=int(data.get("input_channels", 1)),
        output_stride=int(data.get("output_stride", 1)),
        crop_size=crop_size,
        max_instances=data.get("max_instances"),
        max_peaks_per_node=data.get("max_peaks_per_node"),
        max_batch_size=int(data.get("max_batch_size", 1)),
        precision=data.get("precision", "fp32"),
        input_dtype=data.get("input_dtype", "uint8"),
        normalization=data.get("normalization", "0_to_1"),
        n_classes=data.get("n_classes"),
        class_names=data.get("class_names"),
        training_config_embedded=bool(data.get("training_config_embedded", False)),
        training_config_hash=data.get("training_config_hash", ""),
    )

load(path) classmethod

Load from JSON file.

Source code in sleap_nn/export/metadata.py 
108
109
110
111
112
113
@classmethod
def load(cls, path: str | Path) -> "ExportMetadata":
    """Load from JSON file."""
    path = Path(path)
    data = json.loads(path.read_text())
    return cls.from_dict(data)

save(path)

Save to JSON file.

Source code in sleap_nn/export/metadata.py 
103
104
105
106
def save(self, path: str | Path) -> None:
    """Save to JSON file."""
    path = Path(path)
    path.write_text(json.dumps(self.to_dict(), indent=2, sort_keys=True))

to_dict()

Convert to JSON-serializable dict.

Source code in sleap_nn/export/metadata.py 
60
61
62
63
64
65
66
def to_dict(self) -> Dict[str, Any]:
    """Convert to JSON-serializable dict."""
    data = asdict(self)
    data["edge_inds"] = [list(pair) for pair in self.edge_inds]
    if self.crop_size is not None:
        data["crop_size"] = list(self.crop_size)
    return data

ONNXPredictor

Bases: ExportPredictor

ONNX Runtime inference with provider selection.

Methods:

Name Description
__init__

Initialize ONNX predictor with execution providers.
benchmark

Benchmark inference latency and throughput.
predict

Run inference on a batch of images.
Source code in sleap_nn/export/predictors/onnx.py 
13
14
15
16
17
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
80
81
82
83
84
class ONNXPredictor(ExportPredictor):
    """ONNX Runtime inference with provider selection."""

    def __init__(
        self,
        model_path: str,
        device: str = "auto",
        providers: Optional[Iterable[str]] = None,
    ) -> None:
        """Initialize ONNX predictor with execution providers.

        Args:
            model_path: Path to the ONNX model file.
            device: Device for inference ("auto", "cpu", or "cuda").
            providers: ONNX Runtime execution providers. Auto-selected if None.
        """
        try:
            import onnxruntime as ort
        except ImportError as exc:
            raise ImportError(
                "onnxruntime is required for ONNXPredictor. Install with "
                "`pip install onnxruntime` or `onnxruntime-gpu`."
            ) from exc

        # Preload CUDA/cuDNN libraries from pip-installed nvidia packages
        # This is required for onnxruntime-gpu to find the CUDA libraries
        if hasattr(ort, "preload_dlls"):
            ort.preload_dlls()

        self.ort = ort
        if providers is None:
            providers = _select_providers(device, ort.get_available_providers())

        self.session = ort.InferenceSession(model_path, providers=list(providers))
        input_info = self.session.get_inputs()[0]
        self.input_name = input_info.name
        self.input_type = input_info.type
        self.input_dtype = _onnx_type_to_numpy(self.input_type)
        self.output_names = [out.name for out in self.session.get_outputs()]

    def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
        """Run inference on a batch of images."""
        image = _as_numpy(image, expected_dtype=self.input_dtype)
        outputs = self.session.run(None, {self.input_name: image})
        return dict(zip(self.output_names, outputs))

    def benchmark(
        self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
    ) -> Dict[str, float]:
        """Benchmark inference latency and throughput."""
        image = _as_numpy(image, expected_dtype=self.input_dtype)
        for _ in range(n_warmup):
            self.session.run(None, {self.input_name: image})

        times = []
        for _ in range(n_runs):
            start = time.perf_counter()
            self.session.run(None, {self.input_name: image})
            times.append(time.perf_counter() - start)

        times_ms = np.array(times) * 1000.0
        mean_ms = float(times_ms.mean())
        p50_ms = float(np.percentile(times_ms, 50))
        p95_ms = float(np.percentile(times_ms, 95))
        fps = float(1000.0 / mean_ms) if mean_ms > 0 else 0.0

        return {
            "latency_ms_mean": mean_ms,
            "latency_ms_p50": p50_ms,
            "latency_ms_p95": p95_ms,
            "fps": fps,
        }

__init__(model_path, device='auto', providers=None)

Initialize ONNX predictor with execution providers.

Parameters:

Name Type Description Default
model_path str

Path to the ONNX model file.

 required
device str

Device for inference ("auto", "cpu", or "cuda").

 'auto'
providers Optional[Iterable[str]]

ONNX Runtime execution providers. Auto-selected if None.

 None
Source code in sleap_nn/export/predictors/onnx.py 
16
17
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
def __init__(
    self,
    model_path: str,
    device: str = "auto",
    providers: Optional[Iterable[str]] = None,
) -> None:
    """Initialize ONNX predictor with execution providers.

    Args:
        model_path: Path to the ONNX model file.
        device: Device for inference ("auto", "cpu", or "cuda").
        providers: ONNX Runtime execution providers. Auto-selected if None.
    """
    try:
        import onnxruntime as ort
    except ImportError as exc:
        raise ImportError(
            "onnxruntime is required for ONNXPredictor. Install with "
            "`pip install onnxruntime` or `onnxruntime-gpu`."
        ) from exc

    # Preload CUDA/cuDNN libraries from pip-installed nvidia packages
    # This is required for onnxruntime-gpu to find the CUDA libraries
    if hasattr(ort, "preload_dlls"):
        ort.preload_dlls()

    self.ort = ort
    if providers is None:
        providers = _select_providers(device, ort.get_available_providers())

    self.session = ort.InferenceSession(model_path, providers=list(providers))
    input_info = self.session.get_inputs()[0]
    self.input_name = input_info.name
    self.input_type = input_info.type
    self.input_dtype = _onnx_type_to_numpy(self.input_type)
    self.output_names = [out.name for out in self.session.get_outputs()]

benchmark(image, n_warmup=50, n_runs=200)

Benchmark inference latency and throughput.

Source code in sleap_nn/export/predictors/onnx.py 
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
def benchmark(
    self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
) -> Dict[str, float]:
    """Benchmark inference latency and throughput."""
    image = _as_numpy(image, expected_dtype=self.input_dtype)
    for _ in range(n_warmup):
        self.session.run(None, {self.input_name: image})

    times = []
    for _ in range(n_runs):
        start = time.perf_counter()
        self.session.run(None, {self.input_name: image})
        times.append(time.perf_counter() - start)

    times_ms = np.array(times) * 1000.0
    mean_ms = float(times_ms.mean())
    p50_ms = float(np.percentile(times_ms, 50))
    p95_ms = float(np.percentile(times_ms, 95))
    fps = float(1000.0 / mean_ms) if mean_ms > 0 else 0.0

    return {
        "latency_ms_mean": mean_ms,
        "latency_ms_p50": p50_ms,
        "latency_ms_p95": p95_ms,
        "fps": fps,
    }

predict(image)

Run inference on a batch of images.

Source code in sleap_nn/export/predictors/onnx.py 
53
54
55
56
57
def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
    """Run inference on a batch of images."""
    image = _as_numpy(image, expected_dtype=self.input_dtype)
    outputs = self.session.run(None, {self.input_name: image})
    return dict(zip(self.output_names, outputs))

TensorRTPredictor

Bases: ExportPredictor

TensorRT inference for exported models.

This predictor loads a native TensorRT engine file (.trt) and provides inference capabilities using CUDA for high-throughput predictions.

Parameters:

Name Type Description Default
engine_path str | Path

Path to the TensorRT engine file (.trt).

 required
device str

Device to run inference on (only "cuda" supported for TRT).

 'cuda'
Example

predictor = TensorRTPredictor("model.trt") outputs = predictor.predict(images) # uint8 [B, C, H, W]

Methods:

Name Description
__init__

Initialize TensorRT predictor with a serialized engine.
benchmark

Benchmark inference performance.
predict

Run TensorRT inference.
Source code in sleap_nn/export/predictors/tensorrt.py 
 15
 16
 17
 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
 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
class TensorRTPredictor(ExportPredictor):
    """TensorRT inference for exported models.

    This predictor loads a native TensorRT engine file (.trt) and provides
    inference capabilities using CUDA for high-throughput predictions.

    Args:
        engine_path: Path to the TensorRT engine file (.trt).
        device: Device to run inference on (only "cuda" supported for TRT).

    Example:
        >>> predictor = TensorRTPredictor("model.trt")
        >>> outputs = predictor.predict(images)  # uint8 [B, C, H, W]
    """

    def __init__(
        self,
        engine_path: str | Path,
        device: str = "cuda",
    ) -> None:
        """Initialize TensorRT predictor with a serialized engine.

        Args:
            engine_path: Path to the TensorRT engine file.
            device: Device for inference ("cuda" or "auto"). TensorRT requires CUDA.
        """
        import tensorrt as trt

        if device not in ("cuda", "auto"):
            raise ValueError(f"TensorRT only supports CUDA devices, got: {device}")

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

        self.logger = trt.Logger(trt.Logger.WARNING)

        # Load engine
        with open(self.engine_path, "rb") as f:
            self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read())

        if self.engine is None:
            raise RuntimeError(f"Failed to load TensorRT engine: {engine_path}")

        # Create execution context
        self.context = self.engine.create_execution_context()

        # Get input/output info
        self.input_names: List[str] = []
        self.output_names: List[str] = []
        self.input_shapes: Dict[str, tuple] = {}
        self.output_shapes: Dict[str, tuple] = {}

        for i in range(self.engine.num_io_tensors):
            name = self.engine.get_tensor_name(i)
            shape = tuple(self.engine.get_tensor_shape(name))
            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self.input_names.append(name)
                self.input_shapes[name] = shape
            else:
                self.output_names.append(name)
                self.output_shapes[name] = shape

        self.device = torch.device("cuda")

    def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
        """Run TensorRT inference.

        Args:
            image: Input image(s) as numpy array [B, C, H, W] with uint8 dtype.

        Returns:
            Dict mapping output names to numpy arrays.
        """
        import tensorrt as trt

        # Convert input to torch tensor on GPU
        input_tensor = torch.from_numpy(image).to(self.device)

        # Check if engine expects uint8 or float32 and convert if needed
        input_name = self.input_names[0]
        expected_dtype = self.engine.get_tensor_dtype(input_name)
        if expected_dtype == trt.DataType.UINT8:
            # Engine expects uint8 - keep as uint8
            if input_tensor.dtype != torch.uint8:
                input_tensor = input_tensor.to(torch.uint8)
        else:
            # Engine expects float - convert uint8 to float32
            if input_tensor.dtype == torch.uint8:
                input_tensor = input_tensor.to(torch.float32)

        # Ensure contiguous memory
        input_tensor = input_tensor.contiguous()

        # Set input shape for dynamic dimensions
        input_name = self.input_names[0]
        self.context.set_input_shape(input_name, tuple(input_tensor.shape))

        # Allocate output tensors
        outputs: Dict[str, torch.Tensor] = {}
        bindings: Dict[str, int] = {}

        # Set input binding
        bindings[input_name] = input_tensor.data_ptr()

        # Allocate outputs
        for name in self.output_names:
            shape = self.context.get_tensor_shape(name)
            dtype = self._trt_dtype_to_torch(self.engine.get_tensor_dtype(name))
            outputs[name] = torch.empty(tuple(shape), dtype=dtype, device=self.device)
            bindings[name] = outputs[name].data_ptr()

        # Set tensor addresses
        for name, ptr in bindings.items():
            self.context.set_tensor_address(name, ptr)

        # Run inference
        stream = torch.cuda.current_stream().cuda_stream
        success = self.context.execute_async_v3(stream)
        torch.cuda.current_stream().synchronize()

        if not success:
            raise RuntimeError("TensorRT inference failed")

        # Convert outputs to numpy
        return {name: tensor.cpu().numpy() for name, tensor in outputs.items()}

    def benchmark(
        self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
    ) -> Dict[str, float]:
        """Benchmark inference performance.

        Args:
            image: Input image(s) as numpy array [B, C, H, W].
            n_warmup: Number of warmup runs (not timed).
            n_runs: Number of timed runs.

        Returns:
            Dict with timing statistics:
                - mean_ms: Mean inference time in milliseconds
                - std_ms: Standard deviation of inference time
                - min_ms: Minimum inference time
                - max_ms: Maximum inference time
                - fps: Frames per second (based on mean time and batch size)
        """
        batch_size = image.shape[0]

        # Warmup
        for _ in range(n_warmup):
            _ = self.predict(image)

        # Timed runs
        times = []
        for _ in range(n_runs):
            start = time.perf_counter()
            _ = self.predict(image)
            times.append((time.perf_counter() - start) * 1000)

        times_arr = np.array(times)
        mean_ms = float(np.mean(times_arr))

        return {
            "mean_ms": mean_ms,
            "std_ms": float(np.std(times_arr)),
            "min_ms": float(np.min(times_arr)),
            "max_ms": float(np.max(times_arr)),
            "fps": (batch_size * 1000) / mean_ms if mean_ms > 0 else 0.0,
        }

    @staticmethod
    def _trt_dtype_to_torch(trt_dtype):
        """Convert TensorRT dtype to PyTorch dtype."""
        import tensorrt as 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)

__init__(engine_path, device='cuda')

Initialize TensorRT predictor with a serialized engine.

Parameters:

Name Type Description Default
engine_path str | Path

Path to the TensorRT engine file.

 required
device str

Device for inference ("cuda" or "auto"). TensorRT requires CUDA.

 'cuda'
Source code in sleap_nn/export/predictors/tensorrt.py 
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
def __init__(
    self,
    engine_path: str | Path,
    device: str = "cuda",
) -> None:
    """Initialize TensorRT predictor with a serialized engine.

    Args:
        engine_path: Path to the TensorRT engine file.
        device: Device for inference ("cuda" or "auto"). TensorRT requires CUDA.
    """
    import tensorrt as trt

    if device not in ("cuda", "auto"):
        raise ValueError(f"TensorRT only supports CUDA devices, got: {device}")

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

    self.logger = trt.Logger(trt.Logger.WARNING)

    # Load engine
    with open(self.engine_path, "rb") as f:
        self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read())

    if self.engine is None:
        raise RuntimeError(f"Failed to load TensorRT engine: {engine_path}")

    # Create execution context
    self.context = self.engine.create_execution_context()

    # Get input/output info
    self.input_names: List[str] = []
    self.output_names: List[str] = []
    self.input_shapes: Dict[str, tuple] = {}
    self.output_shapes: Dict[str, tuple] = {}

    for i in range(self.engine.num_io_tensors):
        name = self.engine.get_tensor_name(i)
        shape = tuple(self.engine.get_tensor_shape(name))
        if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
            self.input_names.append(name)
            self.input_shapes[name] = shape
        else:
            self.output_names.append(name)
            self.output_shapes[name] = shape

    self.device = torch.device("cuda")

benchmark(image, n_warmup=50, n_runs=200)

Benchmark inference performance.

Parameters:

Name Type Description Default
image ndarray

Input image(s) as numpy array [B, C, H, W].

 required
n_warmup int

Number of warmup runs (not timed).

 50
n_runs int

Number of timed runs.

 200

Returns:

Type Description
Dict[str, float]

Dict with timing statistics: - mean_ms: Mean inference time in milliseconds - std_ms: Standard deviation of inference time - min_ms: Minimum inference time - max_ms: Maximum inference time - fps: Frames per second (based on mean time and batch size)
Source code in sleap_nn/export/predictors/tensorrt.py 
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
def benchmark(
    self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
) -> Dict[str, float]:
    """Benchmark inference performance.

    Args:
        image: Input image(s) as numpy array [B, C, H, W].
        n_warmup: Number of warmup runs (not timed).
        n_runs: Number of timed runs.

    Returns:
        Dict with timing statistics:
            - mean_ms: Mean inference time in milliseconds
            - std_ms: Standard deviation of inference time
            - min_ms: Minimum inference time
            - max_ms: Maximum inference time
            - fps: Frames per second (based on mean time and batch size)
    """
    batch_size = image.shape[0]

    # Warmup
    for _ in range(n_warmup):
        _ = self.predict(image)

    # Timed runs
    times = []
    for _ in range(n_runs):
        start = time.perf_counter()
        _ = self.predict(image)
        times.append((time.perf_counter() - start) * 1000)

    times_arr = np.array(times)
    mean_ms = float(np.mean(times_arr))

    return {
        "mean_ms": mean_ms,
        "std_ms": float(np.std(times_arr)),
        "min_ms": float(np.min(times_arr)),
        "max_ms": float(np.max(times_arr)),
        "fps": (batch_size * 1000) / mean_ms if mean_ms > 0 else 0.0,
    }

predict(image)

Run TensorRT inference.

Parameters:

Name Type Description Default
image ndarray

Input image(s) as numpy array [B, C, H, W] with uint8 dtype.

 required

Returns:

Type Description
Dict[str, ndarray]

Dict mapping output names to numpy arrays.
Source code in sleap_nn/export/predictors/tensorrt.py 
 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
def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
    """Run TensorRT inference.

    Args:
        image: Input image(s) as numpy array [B, C, H, W] with uint8 dtype.

    Returns:
        Dict mapping output names to numpy arrays.
    """
    import tensorrt as trt

    # Convert input to torch tensor on GPU
    input_tensor = torch.from_numpy(image).to(self.device)

    # Check if engine expects uint8 or float32 and convert if needed
    input_name = self.input_names[0]
    expected_dtype = self.engine.get_tensor_dtype(input_name)
    if expected_dtype == trt.DataType.UINT8:
        # Engine expects uint8 - keep as uint8
        if input_tensor.dtype != torch.uint8:
            input_tensor = input_tensor.to(torch.uint8)
    else:
        # Engine expects float - convert uint8 to float32
        if input_tensor.dtype == torch.uint8:
            input_tensor = input_tensor.to(torch.float32)

    # Ensure contiguous memory
    input_tensor = input_tensor.contiguous()

    # Set input shape for dynamic dimensions
    input_name = self.input_names[0]
    self.context.set_input_shape(input_name, tuple(input_tensor.shape))

    # Allocate output tensors
    outputs: Dict[str, torch.Tensor] = {}
    bindings: Dict[str, int] = {}

    # Set input binding
    bindings[input_name] = input_tensor.data_ptr()

    # Allocate outputs
    for name in self.output_names:
        shape = self.context.get_tensor_shape(name)
        dtype = self._trt_dtype_to_torch(self.engine.get_tensor_dtype(name))
        outputs[name] = torch.empty(tuple(shape), dtype=dtype, device=self.device)
        bindings[name] = outputs[name].data_ptr()

    # Set tensor addresses
    for name, ptr in bindings.items():
        self.context.set_tensor_address(name, ptr)

    # Run inference
    stream = torch.cuda.current_stream().cuda_stream
    success = self.context.execute_async_v3(stream)
    torch.cuda.current_stream().synchronize()

    if not success:
        raise RuntimeError("TensorRT inference failed")

    # Convert outputs to numpy
    return {name: tensor.cpu().numpy() for name, tensor in outputs.items()}

build_bottomup_candidate_template(n_nodes, max_peaks_per_node, edge_inds)

Build candidate template matching ONNX wrapper's line_scores ordering.

The ONNX BottomUpONNXWrapper produces line_scores with shape (n_edges, k*k) where for each edge connecting (src_node, dst_node), position i*k + j corresponds to: - src peak flat index: src_node * k + i - dst peak flat index: dst_node * k + j

This function builds edge_inds and edge_peak_inds tensors that match this exact ordering, so that line_scores_flat[idx] corresponds to edge_peak_inds[idx].

Parameters:

Name Type Description Default
n_nodes int

Number of nodes in the skeleton.

 required
max_peaks_per_node int

Maximum peaks per node (k) used during export.

 required
edge_inds List[Tuple[int, int]]

List of (src_node, dst_node) tuples defining skeleton edges.

 required

Returns:

Type Description
Tuple['torch.Tensor', 'torch.Tensor', 'torch.Tensor']

Tuple of (peak_channel_inds, edge_inds_tensor, edge_peak_inds_tensor): - peak_channel_inds: (n_nodes * k,) tensor mapping flat peak index to node - edge_inds_tensor: (n_edges * k * k,) tensor of edge indices for each candidate - edge_peak_inds_tensor: (n_edges * k * k, 2) tensor of (src, dst) peak indices
Example

from sleap_nn.export.utils import build_bottomup_candidate_template peak_ch, edge_inds, edge_peaks = build_bottomup_candidate_template( ... n_nodes=15, max_peaks_per_node=20, edge_inds=[(1, 2), (1, 5)] ... )

Use with ONNX output:

line_scores_flat = line_scores.reshape(-1) valid_scores = line_scores_flat[valid_mask] valid_edge_peaks = edge_peaks[valid_mask]

Note

This function is necessary because get_connection_candidates() in sleap_nn.inference.paf_grouping uses unstable argsort, which shuffles peak indices within each node and breaks alignment with ONNX output ordering.

Source code in sleap_nn/export/utils.py 
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
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
def build_bottomup_candidate_template(
    n_nodes: int, max_peaks_per_node: int, edge_inds: List[Tuple[int, int]]
) -> Tuple["torch.Tensor", "torch.Tensor", "torch.Tensor"]:
    """Build candidate template matching ONNX wrapper's line_scores ordering.

    The ONNX BottomUpONNXWrapper produces line_scores with shape (n_edges, k*k) where
    for each edge connecting (src_node, dst_node), position i*k + j corresponds to:
    - src peak flat index: src_node * k + i
    - dst peak flat index: dst_node * k + j

    This function builds edge_inds and edge_peak_inds tensors that match this exact
    ordering, so that line_scores_flat[idx] corresponds to edge_peak_inds[idx].

    Args:
        n_nodes: Number of nodes in the skeleton.
        max_peaks_per_node: Maximum peaks per node (k) used during export.
        edge_inds: List of (src_node, dst_node) tuples defining skeleton edges.

    Returns:
        Tuple of (peak_channel_inds, edge_inds_tensor, edge_peak_inds_tensor):
        - peak_channel_inds: (n_nodes * k,) tensor mapping flat peak index to node
        - edge_inds_tensor: (n_edges * k * k,) tensor of edge indices for each candidate
        - edge_peak_inds_tensor: (n_edges * k * k, 2) tensor of (src, dst) peak indices

    Example:
        >>> from sleap_nn.export.utils import build_bottomup_candidate_template
        >>> peak_ch, edge_inds, edge_peaks = build_bottomup_candidate_template(
        ...     n_nodes=15, max_peaks_per_node=20, edge_inds=[(1, 2), (1, 5)]
        ... )
        >>> # Use with ONNX output:
        >>> line_scores_flat = line_scores.reshape(-1)
        >>> valid_scores = line_scores_flat[valid_mask]
        >>> valid_edge_peaks = edge_peaks[valid_mask]

    Note:
        This function is necessary because `get_connection_candidates()` in
        `sleap_nn.inference.paf_grouping` uses unstable argsort, which shuffles
        peak indices within each node and breaks alignment with ONNX output ordering.
    """
    import torch

    k = max_peaks_per_node
    n_edges = len(edge_inds)

    # peak_channel_inds: [0,0,...0, 1,1,...1, ...] (k times each)
    peak_channel_inds = torch.arange(n_nodes, dtype=torch.int32).repeat_interleave(k)

    edge_inds_list = []
    edge_peak_inds_list = []

    for edge_idx, (src_node, dst_node) in enumerate(edge_inds):
        # Build k*k candidate pairs in row-major order (i*k + j)
        # src indices: [src_node*k + 0, src_node*k + 0, ..., src_node*k + 1, ...]
        # dst indices: [dst_node*k + 0, dst_node*k + 1, ..., dst_node*k + 0, ...]
        src_base = src_node * k
        dst_base = dst_node * k

        src_indices = torch.arange(k, dtype=torch.int32).repeat_interleave(k) + src_base
        dst_indices = torch.arange(k, dtype=torch.int32).repeat(k) + dst_base

        edge_inds_list.append(torch.full((k * k,), edge_idx, dtype=torch.int32))
        edge_peak_inds_list.append(torch.stack([src_indices, dst_indices], dim=1))

    if edge_inds_list:
        edge_inds_tensor = torch.cat(edge_inds_list)
        edge_peak_inds_tensor = torch.cat(edge_peak_inds_list)
    else:
        edge_inds_tensor = torch.empty((0,), dtype=torch.int32)
        edge_peak_inds_tensor = torch.empty((0, 2), dtype=torch.int32)

    return peak_channel_inds, edge_inds_tensor, edge_peak_inds_tensor

export_model(model, save_path, fmt='onnx', input_shape=(1, 1, 512, 512), opset_version=17, output_names=None, verify=True, **kwargs)

Export a model to the requested format.

Source code in sleap_nn/export/exporters/__init__.py 
14
15
16
17
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
def export_model(
    model: torch.nn.Module,
    save_path: str | Path,
    fmt: str = "onnx",
    input_shape: Iterable[int] = (1, 1, 512, 512),
    opset_version: int = 17,
    output_names: Optional[list] = None,
    verify: bool = True,
    **kwargs,
) -> Path:
    """Export a model to the requested format."""
    fmt = fmt.lower()
    if fmt == "onnx":
        return export_to_onnx(
            model,
            save_path,
            input_shape=input_shape,
            opset_version=opset_version,
            output_names=output_names,
            verify=verify,
        )
    if fmt == "tensorrt":
        return export_to_tensorrt(model, save_path, input_shape=input_shape, **kwargs)
    if fmt == "both":
        export_to_onnx(
            model,
            save_path,
            input_shape=input_shape,
            opset_version=opset_version,
            output_names=output_names,
            verify=verify,
        )
        return export_to_tensorrt(model, save_path, input_shape=input_shape, **kwargs)

    raise ValueError(f"Unknown export format: {fmt}")

export_to_onnx(model, save_path, input_shape=(1, 1, 512, 512), input_dtype=torch.uint8, opset_version=17, dynamic_axes=None, input_names=None, output_names=None, do_constant_folding=True, verify=True)

Export a PyTorch model to ONNX.

Source code in sleap_nn/export/exporters/onnx_exporter.py 
11
12
13
14
15
16
17
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
def export_to_onnx(
    model: torch.nn.Module,
    save_path: str | Path,
    input_shape: Iterable[int] = (1, 1, 512, 512),
    input_dtype: torch.dtype = torch.uint8,
    opset_version: int = 17,
    dynamic_axes: Optional[Dict[str, Dict[int, str]]] = None,
    input_names: Optional[List[str]] = None,
    output_names: Optional[List[str]] = None,
    do_constant_folding: bool = True,
    verify: bool = True,
) -> Path:
    """Export a PyTorch model to ONNX."""
    save_path = Path(save_path)
    model.eval()

    if input_names is None:
        input_names = ["image"]
    if dynamic_axes is None:
        dynamic_axes = {"image": {0: "batch", 2: "height", 3: "width"}}

    device = None
    try:
        device = next(model.parameters()).device
    except StopIteration:
        device = torch.device("cpu")

    if input_dtype.is_floating_point:
        dummy_input = torch.randn(*input_shape, device=device, dtype=input_dtype)
    else:
        dummy_input = torch.randint(
            0, 256, input_shape, device=device, dtype=input_dtype
        )

    if output_names is None:
        with torch.no_grad():
            test_out = model(dummy_input)
        output_names = _infer_output_names(test_out)

    torch.onnx.export(
        model,
        dummy_input,
        save_path.as_posix(),
        opset_version=opset_version,
        input_names=input_names,
        output_names=output_names,
        dynamic_axes=dynamic_axes,
        do_constant_folding=do_constant_folding,
        dynamo=False,
    )

    if verify:
        _verify_onnx(save_path)

    return save_path

export_to_tensorrt(model, save_path, input_shape=(1, 1, 512, 512), input_dtype=torch.uint8, precision='fp16', min_shape=None, opt_shape=None, max_shape=None, workspace_size=2 << 30, method='onnx', verbose=True)

Export a PyTorch model to TensorRT format.

This function supports multiple compilation methods: - "onnx": Exports to ONNX first, then compiles with TensorRT (most reliable) - "jit": Uses torch.jit.trace + torch_tensorrt.compile (alternative)

Parameters:

Name Type Description Default
model Module

The PyTorch model to export (typically an ONNX wrapper).

 required
save_path str | Path

Path to save the TensorRT engine (.trt file).

 required
input_shape Tuple[int, int, int, int]

(B, C, H, W) optimal input tensor shape.

 (1, 1, 512, 512)
input_dtype dtype

Input tensor dtype (torch.uint8 or torch.float32).

 uint8
precision str

Model precision - "fp32" or "fp16".

 'fp16'
min_shape Optional[Tuple[int, int, int, int]]

Minimum input shape for dynamic shapes (default: batch=1, H/W halved).

 None
opt_shape Optional[Tuple[int, int, int, int]]

Optimal input shape (default: same as input_shape).

 None
max_shape Optional[Tuple[int, int, int, int]]

Maximum input shape (default: batch=16, H/W doubled).

 None
workspace_size int

TensorRT workspace size in bytes (default 2GB).

 2 << 30
method str

Compilation method - "onnx" or "jit".

 'onnx'
verbose bool

Print export info.

 True

Returns:

Type Description
Path

Path to the exported TensorRT engine.
Note

TensorRT models are NOT cross-platform. The exported model will only work on the same GPU architecture and TensorRT version used for export.

Source code in sleap_nn/export/exporters/tensorrt_exporter.py 
 12
 13
 14
 15
 16
 17
 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
 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
def export_to_tensorrt(
    model: nn.Module,
    save_path: str | Path,
    input_shape: Tuple[int, int, int, int] = (1, 1, 512, 512),
    input_dtype: torch.dtype = torch.uint8,
    precision: str = "fp16",
    min_shape: Optional[Tuple[int, int, int, int]] = None,
    opt_shape: Optional[Tuple[int, int, int, int]] = None,
    max_shape: Optional[Tuple[int, int, int, int]] = None,
    workspace_size: int = 2 << 30,  # 2GB default
    method: str = "onnx",
    verbose: bool = True,
) -> Path:
    """Export a PyTorch model to TensorRT format.

    This function supports multiple compilation methods:
    - "onnx": Exports to ONNX first, then compiles with TensorRT (most reliable)
    - "jit": Uses torch.jit.trace + torch_tensorrt.compile (alternative)

    Args:
        model: The PyTorch model to export (typically an ONNX wrapper).
        save_path: Path to save the TensorRT engine (.trt file).
        input_shape: (B, C, H, W) optimal input tensor shape.
        input_dtype: Input tensor dtype (torch.uint8 or torch.float32).
        precision: Model precision - "fp32" or "fp16".
        min_shape: Minimum input shape for dynamic shapes (default: batch=1, H/W halved).
        opt_shape: Optimal input shape (default: same as input_shape).
        max_shape: Maximum input shape (default: batch=16, H/W doubled).
        workspace_size: TensorRT workspace size in bytes (default 2GB).
        method: Compilation method - "onnx" or "jit".
        verbose: Print export info.

    Returns:
        Path to the exported TensorRT engine.

    Note:
        TensorRT models are NOT cross-platform. The exported model will only
        work on the same GPU architecture and TensorRT version used for export.
    """
    import tensorrt as trt

    model.eval()
    device = next(model.parameters()).device

    save_path = Path(save_path)
    if not save_path.suffix:
        save_path = save_path.with_suffix(".trt")

    B, C, H, W = input_shape

    if min_shape is None:
        min_shape = (1, C, H // 2, W // 2)
    if opt_shape is None:
        opt_shape = input_shape
    if max_shape is None:
        max_shape = (min(16, B * 4), C, H * 2, W * 2)

    if verbose:
        print(f"Exporting model to TensorRT...")
        print(f"  Input shape: {input_shape}")
        print(f"  Min/Opt/Max: {min_shape} / {opt_shape} / {max_shape}")
        print(f"  Precision: {precision}")
        print(f"  Workspace: {workspace_size / 1e9:.1f} GB")
        print(f"  Method: {method}")

    if method == "onnx":
        return _export_tensorrt_onnx(
            model,
            save_path,
            input_shape,
            input_dtype,
            min_shape,
            opt_shape,
            max_shape,
            precision,
            workspace_size,
            verbose,
        )
    elif method == "jit":
        return _export_tensorrt_jit(
            model,
            save_path,
            input_shape,
            input_dtype,
            min_shape,
            opt_shape,
            max_shape,
            precision,
            workspace_size,
            verbose,
        )
    else:
        raise ValueError(f"Unknown method: {method}. Use 'onnx' or 'jit'.")

load_exported_model(model_path, runtime='auto', device='auto', **kwargs)

Load an exported model and return a predictor instance.

Source code in sleap_nn/export/predictors/__init__.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def load_exported_model(
    model_path: str | Path,
    runtime: str = "auto",
    device: str = "auto",
    **kwargs,
) -> ExportPredictor:
    """Load an exported model and return a predictor instance."""
    if runtime == "auto":
        runtime = detect_runtime(model_path)

    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"

    if runtime == "onnx":
        return ONNXPredictor(str(model_path), device=device, **kwargs)
    if runtime == "tensorrt":
        return TensorRTPredictor(str(model_path), device=device, **kwargs)

    raise ValueError(f"Unknown runtime: {runtime}")