Skip to content

onnx

sleap_nn.export.predictors.onnx

ONNX Runtime predictor.

Classes:

Name Description
ONNXPredictor

ONNX Runtime inference with provider selection.

ONNXPredictor

Bases: ExportPredictor

ONNX Runtime inference with provider selection.

Methods:

Name Description
__init__

Initialize ONNX predictor with execution providers.
benchmark

Benchmark inference latency and throughput.
predict

Run inference on a batch of images.
Source code in sleap_nn/export/predictors/onnx.py 
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
class ONNXPredictor(ExportPredictor):
    """ONNX Runtime inference with provider selection."""

    def __init__(
        self,
        model_path: str,
        device: str = "auto",
        providers: Optional[Iterable[str]] = None,
    ) -> None:
        """Initialize ONNX predictor with execution providers.

        Args:
            model_path: Path to the ONNX model file.
            device: Device for inference ("auto", "cpu", or "cuda").
            providers: ONNX Runtime execution providers. Auto-selected if None.
        """
        try:
            import onnxruntime as ort
        except ImportError as exc:
            raise ImportError(
                "onnxruntime is required for ONNXPredictor. Install with "
                "`pip install onnxruntime` or `onnxruntime-gpu`."
            ) from exc

        # Preload CUDA/cuDNN libraries from pip-installed nvidia packages
        # This is required for onnxruntime-gpu to find the CUDA libraries
        if hasattr(ort, "preload_dlls"):
            ort.preload_dlls()

        self.ort = ort
        if providers is None:
            providers = _select_providers(device, ort.get_available_providers())

        self.session = ort.InferenceSession(model_path, providers=list(providers))
        input_info = self.session.get_inputs()[0]
        self.input_name = input_info.name
        self.input_type = input_info.type
        self.input_dtype = _onnx_type_to_numpy(self.input_type)
        self.output_names = [out.name for out in self.session.get_outputs()]

    def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
        """Run inference on a batch of images."""
        image = _as_numpy(image, expected_dtype=self.input_dtype)
        outputs = self.session.run(None, {self.input_name: image})
        return dict(zip(self.output_names, outputs))

    def benchmark(
        self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
    ) -> Dict[str, float]:
        """Benchmark inference latency and throughput."""
        image = _as_numpy(image, expected_dtype=self.input_dtype)
        for _ in range(n_warmup):
            self.session.run(None, {self.input_name: image})

        times = []
        for _ in range(n_runs):
            start = time.perf_counter()
            self.session.run(None, {self.input_name: image})
            times.append(time.perf_counter() - start)

        times_ms = np.array(times) * 1000.0
        mean_ms = float(times_ms.mean())
        p50_ms = float(np.percentile(times_ms, 50))
        p95_ms = float(np.percentile(times_ms, 95))
        fps = float(1000.0 / mean_ms) if mean_ms > 0 else 0.0

        return {
            "latency_ms_mean": mean_ms,
            "latency_ms_p50": p50_ms,
            "latency_ms_p95": p95_ms,
            "fps": fps,
        }

__init__(model_path, device='auto', providers=None)

Initialize ONNX predictor with execution providers.

Parameters:

Name Type Description Default
model_path str

Path to the ONNX model file.

 required
device str

Device for inference ("auto", "cpu", or "cuda").

 'auto'
providers Optional[Iterable[str]]

ONNX Runtime execution providers. Auto-selected if None.

 None
Source code in sleap_nn/export/predictors/onnx.py 
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
def __init__(
    self,
    model_path: str,
    device: str = "auto",
    providers: Optional[Iterable[str]] = None,
) -> None:
    """Initialize ONNX predictor with execution providers.

    Args:
        model_path: Path to the ONNX model file.
        device: Device for inference ("auto", "cpu", or "cuda").
        providers: ONNX Runtime execution providers. Auto-selected if None.
    """
    try:
        import onnxruntime as ort
    except ImportError as exc:
        raise ImportError(
            "onnxruntime is required for ONNXPredictor. Install with "
            "`pip install onnxruntime` or `onnxruntime-gpu`."
        ) from exc

    # Preload CUDA/cuDNN libraries from pip-installed nvidia packages
    # This is required for onnxruntime-gpu to find the CUDA libraries
    if hasattr(ort, "preload_dlls"):
        ort.preload_dlls()

    self.ort = ort
    if providers is None:
        providers = _select_providers(device, ort.get_available_providers())

    self.session = ort.InferenceSession(model_path, providers=list(providers))
    input_info = self.session.get_inputs()[0]
    self.input_name = input_info.name
    self.input_type = input_info.type
    self.input_dtype = _onnx_type_to_numpy(self.input_type)
    self.output_names = [out.name for out in self.session.get_outputs()]

benchmark(image, n_warmup=50, n_runs=200)

Benchmark inference latency and throughput.

Source code in sleap_nn/export/predictors/onnx.py 
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
def benchmark(
    self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
) -> Dict[str, float]:
    """Benchmark inference latency and throughput."""
    image = _as_numpy(image, expected_dtype=self.input_dtype)
    for _ in range(n_warmup):
        self.session.run(None, {self.input_name: image})

    times = []
    for _ in range(n_runs):
        start = time.perf_counter()
        self.session.run(None, {self.input_name: image})
        times.append(time.perf_counter() - start)

    times_ms = np.array(times) * 1000.0
    mean_ms = float(times_ms.mean())
    p50_ms = float(np.percentile(times_ms, 50))
    p95_ms = float(np.percentile(times_ms, 95))
    fps = float(1000.0 / mean_ms) if mean_ms > 0 else 0.0

    return {
        "latency_ms_mean": mean_ms,
        "latency_ms_p50": p50_ms,
        "latency_ms_p95": p95_ms,
        "fps": fps,
    }

predict(image)

Run inference on a batch of images.

Source code in sleap_nn/export/predictors/onnx.py 
53
54
55
56
57
def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
    """Run inference on a batch of images."""
    image = _as_numpy(image, expected_dtype=self.input_dtype)
    outputs = self.session.run(None, {self.input_name: image})
    return dict(zip(self.output_names, outputs))