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)

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