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.
utils	Utilities for export workflows.
wrappers	ONNX/TensorRT export wrappers.

Classes:

Name	Description
ExportMetadata	Metadata embedded or saved alongside 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.

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

@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"
    peak_threshold: Optional[float] = None

    # 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

    # Centroid/top-down anchor point, and how the centroid is derived (#586).
    # ``centroid_method`` is what `resolve_centroid_method` returned at export
    # time; ``None`` in older exports, which then infer it from ``anchor_part``.
    anchor_part: Optional[str] = None
    centroid_method: Optional[str] = None

    # Embedding (re-ID) model
    embedding_dim: Optional[int] = None
    normalize: Optional[bool] = None
    backbone_source: Optional[str] = None
    # ``burn_in``/``background_fill`` record whether the trained embedder masked the
    # crop (foreground-only standardize). The single-input ONNX wrapper always does a
    # maskless whole-crop standardize, so a ``burn_in=True`` model's exported embeddings
    # DIVERGE from native masked inference — these fields make that detectable.
    burn_in: Optional[bool] = None
    background_fill: Optional[str] = None
    # The crop geometry the embedder was trained on, which a runtime must
    # reproduce to feed it matching crops: how the crop is centered on a detection
    # and the frame size-matching applied before cropping (``None`` = none).
    crop_centering: Optional[str] = None
    max_height: Optional[int] = None
    max_width: Optional[int] = 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"),
            peak_threshold=data.get("peak_threshold"),
            anchor_part=data.get("anchor_part"),
            embedding_dim=data.get("embedding_dim"),
            normalize=data.get("normalize"),
            backbone_source=data.get("backbone_source"),
            burn_in=data.get("burn_in"),
            background_fill=data.get("background_fill"),
            crop_centering=data.get("crop_centering"),
            max_height=data.get("max_height"),
            max_width=data.get("max_width"),
            centroid_method=data.get("centroid_method"),
            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

@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

@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"),
        peak_threshold=data.get("peak_threshold"),
        anchor_part=data.get("anchor_part"),
        embedding_dim=data.get("embedding_dim"),
        normalize=data.get("normalize"),
        backbone_source=data.get("backbone_source"),
        burn_in=data.get("burn_in"),
        background_fill=data.get("background_fill"),
        crop_centering=data.get("crop_centering"),
        max_height=data.get("max_height"),
        max_width=data.get("max_width"),
        centroid_method=data.get("centroid_method"),
        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

@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

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

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

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

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

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, numerical_check=False, numerical_atol=0.001, numerical_rtol=0.001, static_hw=None)

Export a PyTorch model to ONNX.

Parameters:

Name	Type	Description	Default
model	Module	The PyTorch module to export.	required
save_path	str | Path	Destination path for the .onnx file.	required
input_shape	Iterable[int]	Shape of the dummy input used to trace the graph.	(1, 1, 512, 512)
input_dtype	dtype	Dtype of the dummy input (uint8 matches the inference wrappers, which normalize in-graph).	uint8
opset_version	int	ONNX opset for the (default) TorchScript exporter.	17
dynamic_axes	Optional[Dict[str, Dict[int, str]]]	Dynamic-axis spec; defaults to batch/height/width dynamic on the image input.	None
input_names	Optional[List[str]]	ONNX input names; defaults to ["image"].	None
output_names	Optional[List[str]]	ONNX output names; inferred from a reference forward if None.	None
do_constant_folding	bool	Whether to constant-fold during export.	True
verify	bool	If True, run the structural onnx.checker on the result.	True
numerical_check	bool	If True, additionally run the exported graph through onnxruntime on the export dummy input and assert output parity against PyTorch (atol/rtol below). _verify_onnx alone is structural only (onnx.checker); a numerical check catches graphs that are valid but wrong — a known failure mode for transformer backbones (e.g. Swin) whose ops trace but disagree numerically. Requires onnxruntime (the export extra); degrades to a warning if it is unavailable.	False
numerical_atol	float	Absolute tolerance for the numerical parity check.	0.001
numerical_rtol	float	Relative tolerance for the numerical parity check.	0.001
static_hw	Optional[bool]	If True, export the input's height and width as static (the traced input_shape); only the batch stays dynamic. None (default) does so when the model contains a rotary position embedding (has_rope), whose grid tracing freezes. Pass False when the RoPE module only ever sees a fixed size anyway -- e.g. the instance model inside a top-down wrapper, which receives fixed-size crops.	None

Source code in sleap_nn/export/exporters/onnx_exporter.py

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,
    numerical_check: bool = False,
    numerical_atol: float = 1e-3,
    numerical_rtol: float = 1e-3,
    static_hw: Optional[bool] = None,
) -> Path:
    """Export a PyTorch model to ONNX.

    Args:
        model: The PyTorch module to export.
        save_path: Destination path for the ``.onnx`` file.
        input_shape: Shape of the dummy input used to trace the graph.
        input_dtype: Dtype of the dummy input (``uint8`` matches the inference
            wrappers, which normalize in-graph).
        opset_version: ONNX opset for the (default) TorchScript exporter.
        dynamic_axes: Dynamic-axis spec; defaults to batch/height/width dynamic on
            the ``image`` input.
        input_names: ONNX input names; defaults to ``["image"]``.
        output_names: ONNX output names; inferred from a reference forward if
            ``None``.
        do_constant_folding: Whether to constant-fold during export.
        verify: If ``True``, run the structural ``onnx.checker`` on the result.
        numerical_check: If ``True``, additionally run the exported graph through
            onnxruntime on the export dummy input and assert output parity against
            PyTorch (atol/rtol below). ``_verify_onnx`` alone is structural only
            (``onnx.checker``); a numerical check catches graphs that are valid but
            wrong — a known failure mode for transformer backbones (e.g. Swin) whose
            ops trace but disagree numerically. Requires onnxruntime (the ``export``
            extra); degrades to a warning if it is unavailable.
        numerical_atol: Absolute tolerance for the numerical parity check.
        numerical_rtol: Relative tolerance for the numerical parity check.
        static_hw: If ``True``, export the input's height and width as static
            (the traced ``input_shape``); only the batch stays dynamic. ``None``
            (default) does so when the model contains a rotary position embedding
            (`has_rope`), whose grid tracing freezes. Pass ``False`` when the RoPE
            module only ever sees a fixed size anyway -- e.g. the instance model
            inside a top-down wrapper, which receives fixed-size crops.
    """
    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"}}
    rope = has_rope(model)
    if static_hw is None:
        static_hw = rope
    if static_hw:
        # A RoPE grid is only correct at the traced size. Export H/W static so
        # onnxruntime / TensorRT reject other sizes instead of returning silently
        # wrong outputs.
        dynamic_axes = _without_spatial_axes(dynamic_axes, input_names[0])
        reason = (
            "The backbone uses rotary position embeddings, which trace to a fixed "
            "token grid"
            if rope
            else "static_hw=True"
        )
        logger.warning(
            f"{reason}: exporting with static height/width "
            f"{tuple(input_shape)[2:]}; feed the ONNX model inputs of that size."
        )

    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
        )

    # A reference forward is needed to infer output names and/or for parity.
    test_out = None
    if output_names is None or numerical_check:
        with torch.no_grad():
            test_out = model(dummy_input)
    if output_names is None:
        output_names = _infer_output_names(test_out)

    common = dict(
        input_names=input_names,
        output_names=output_names,
        dynamic_axes=dynamic_axes,
    )

    def _export_torchscript(opset: int) -> None:
        # Default to the legacy TorchScript exporter: it is fast and exports every
        # current wrapper (including the multi-stage top-down / multi-class ones that
        # the torch.export-based exporter cannot trace).
        torch.onnx.export(
            model,
            dummy_input,
            save_path.as_posix(),
            opset_version=opset,
            do_constant_folding=do_constant_folding,
            dynamo=False,
            **common,
        )

    torch.onnx.register_custom_op_symbolic(
        "aten::_upsample_bicubic2d_aa",
        _upsample_bicubic2d_aa_symbolic,
        _ANTIALIAS_OPSET,
    )
    try:
        _export_torchscript(opset_version)
    except torch.onnx.errors.UnsupportedOperatorError as err:
        try:
            if (
                opset_version >= _ANTIALIAS_OPSET
                or "_upsample_bicubic2d_aa" not in str(err)
            ):
                raise
            # A plain-ViT backbone's positional-embedding resample (antialiased
            # bicubic) lowers at opset 18 via the symbolic registered above. Stay on
            # the TorchScript exporter for it: it traces HF's shape arithmetic
            # dynamically, whereas torch.export pins the token grid (and, through
            # the [1, 3, 1, 1] mean/std broadcast, the batch) to the dummy input.
            _export_torchscript(_ANTIALIAS_OPSET)
        except torch.onnx.errors.UnsupportedOperatorError:
            # A few ops have no symbolic in the legacy exporter — notably the
            # antialiased resize (``aten::_upsample_bilinear2d_aa``) that
            # single-instance / other downscaling wrappers use to match the PyTorch
            # inference resize. The torch.export-based exporter supports them, so
            # fall back to it for those models (it needs opset >= 18 for the
            # antialias Resize attribute).
            torch.onnx.export(
                model,
                dummy_input,
                save_path.as_posix(),
                opset_version=max(opset_version, _ANTIALIAS_OPSET),
                dynamo=True,
                **common,
            )

    if verify:
        _verify_onnx(save_path)

    if numerical_check:
        _verify_onnx_numerical(
            save_path,
            dummy_input,
            test_out,
            input_names[0],
            output_names,
            atol=numerical_atol,
            rtol=numerical_rtol,
        )

    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, static_hw=None)

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
static_hw	Optional[bool]	Whether the engine takes only the traced height/width (a RoPE backbone's grid is fixed at trace time). None detects RoPE in model (has_rope); pass False for a wrapper whose RoPE model only sees fixed-size crops (a top-down instance model). For the "onnx" method an existing ONNX file's own static dims win.	None

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

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,
    static_hw: Optional[bool] = None,
) -> 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.
        static_hw: Whether the engine takes only the traced height/width (a RoPE
            backbone's grid is fixed at trace time). ``None`` detects RoPE in
            ``model`` (`has_rope`); pass ``False`` for a wrapper whose RoPE model
            only sees fixed-size crops (a top-down instance model). For the
            ``"onnx"`` method an existing ONNX file's own static dims win.

    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 static_hw is None:
        static_hw = has_rope(model)
    if method == "jit" and static_hw:
        # The traced RoPE grid is only correct at the traced size (the ONNX route
        # gets the same restriction from the static ONNX input instead).
        pinned = tuple((s[0], s[1], H, W) for s in (min_shape, opt_shape, max_shape))
        if pinned != (tuple(min_shape), tuple(opt_shape), tuple(max_shape)):
            logger.warning(
                "The backbone uses rotary position embeddings, which trace to a "
                f"fixed token grid: TensorRT profile height/width pinned to {H}x{W}."
            )
        min_shape, opt_shape, max_shape = pinned

    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,
            static_hw,
        )
    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'.")