onnx_exporter

sleap_nn.export.exporters.onnx_exporter

ONNX export utilities.

Functions:

Name	Description
export_to_onnx	Export a PyTorch model to ONNX.
has_rope	Whether a model contains a rotary position embedding module.

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

has_rope(model)

Whether a model contains a rotary position embedding module.

Rotary position embeddings (timm DINOv3 / EVA02 / PE) compute their grid from Python floats of the input size, which tracing freezes, so a traced graph is only correct at the size it was traced at.

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

def has_rope(model: torch.nn.Module) -> bool:
    """Whether a model contains a rotary position embedding module.

    Rotary position embeddings (timm DINOv3 / EVA02 / PE) compute their grid from
    Python floats of the input size, which tracing freezes, so a traced graph is
    only correct at the size it was traced at.
    """
    return any(
        "rope" in name or "rotary" in name
        for name in (type(m).__name__.lower() for m in model.modules())
    )