Skip to content

onnx_exporter

sleap_nn.export.exporters.onnx_exporter

ONNX export utilities.

Functions:

Name Description
export_to_onnx

Export a PyTorch model to ONNX.

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 
 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
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
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