exporters

sleap_nn.export.exporters

Exporters for serialized model formats.

Modules:

Name	Description
onnx_exporter	ONNX export utilities.
tensorrt_exporter	TensorRT export utilities.

Functions:

Name	Description
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.

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)

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

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

    # 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,
    )
    try:
        # 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_version,
            do_constant_folding=do_constant_folding,
            dynamo=False,
            **common,
        )
    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, 18),
            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)

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

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