Runtime (ONNX Runtime, CPU / GPU)

Runtime (ONNX Runtime, CPU / GPU)#

ONNX Runtime layer shared by every model in the package (inference only).

OnnxModel wraps one onnxruntime.InferenceSession:

  • device: 'auto' (CUDA if it really works, else CPU), 'cpu' or 'cuda'. 'cuda' never falls back silently: if the CUDA execution provider is missing or cannot be initialised, it raises RuntimeError with what to install. (ONNX Runtime itself silently creates a CPU session when the CUDA libraries fail to load; this wrapper checks the providers the session actually got.)

  • threads: number of CPU threads for one inference call (intra_op_num_threads); None = ONNX Runtime’s default (one per physical core).

  • batching: OnnxModel.run() splits the inputs along the first (batch) axis into chunks of batch_size, casts and checks each chunk just before its inference call, and concatenates the outputs along each output’s batch axis (checked on every chunk, see output_batch_axes). This bounds the memory of each inference call (activations and the cast copy of the input); the outputs of the whole batch are of course all kept.

  • integrity: an optional SHA-256 of the model file is verified before loading, so a truncated or swapped model file fails loudly instead of producing wrong numbers.

  • inputs: every input is cast to the dtype the model declares and must be finite (NaN or inf reaching a network silently corrupts its output; the detectors handle gaps before).

GPU installation (NVIDIA)#

There is deliberately no [gpu] extra. onnxruntime (CPU) and onnxruntime-gpu install the same Python module onnxruntime: with both installed, the CPU build typically wins without any error (pip check passes), and uninstalling one breaks the other. Swap the package by hand:

pip install brainmaze-eeg-models
pip uninstall -y onnxruntime
pip install "onnxruntime-gpu[cuda,cudnn]"     # CUDA + cuDNN as pip wheels
python -c "import brainmaze_eeg_models as bm; print(bm.check_gpu())"

check_gpu() is the real check: it creates a CUDA session and runs a tiny inference on the GPU, and raises RuntimeError with the reason (and what to install) if that does not work. A CUDAExecutionProvider in onnxruntime.get_available_providers() only shows that the GPU build is installed: a GPU build without usable CUDA libraries lists it too, and then every session falls back to the CPU. Any model created with device='cuda' (e.g. SpindleDetector('eeg', device='cuda')) also raises if its CUDA session cannot be created. device_report() lists what is installed.

Upgrades: pip install -U brainmaze-eeg-models re-installs the CPU onnxruntime (it is a dependency), which overwrites the GPU module again. Upgrade with pip install -U --no-deps brainmaze-eeg-models (or redo the swap afterwards). With device='auto' a clobbered install is reported by a RuntimeWarning.

CUDA versions: onnxruntime-gpu >= 1.27 on PyPI is built for CUDA 13.0 + cuDNN 9 (needs a driver that supports CUDA 13). For a CUDA 12 system (e.g. an older cluster driver) the simplest route is the last CUDA 12 builds on PyPI, 1.24-1.26 (Python >= 3.11), whose [cuda,cudnn] extras install the -cu12 CUDA and cuDNN wheels:

pip uninstall -y onnxruntime onnxruntime-gpu
pip install "onnxruntime-gpu[cuda,cudnn]<1.27"

For a newer CUDA 12 build use Microsoft’s CUDA 12 feed, whose builds need CUDA >= 12.8 (1.27-1.29, Python 3.11-3.14 on 2026-10-02; pick the newest version listed there):

pip uninstall -y onnxruntime onnxruntime-gpu
pip install --no-deps onnxruntime-gpu==1.29.0         --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/
pip install nvidia-cuda-runtime-cu12 nvidia-cudnn-cu12 nvidia-cublas-cu12         nvidia-cufft-cu12 nvidia-curand-cu12 nvidia-cuda-nvrtc-cu12   # or a system CUDA >= 12.8 + cuDNN 9

(--no-deps: the feed also mirrors other packages, in old versions (without it pip would install e.g. numpy 2.1.2 and protobuf 5.28.3 from the feed); the dependencies are already installed with the CPU package. With --index-url pip does not look at PyPI at all, so it cannot pick the CUDA 13 build from there.) Python 3.10: the newest onnxruntime-gpu with Python 3.10 wheels is 1.23.2, a CUDA 12 build, so the first recipe gives CUDA 12 (its [cuda,cudnn] extras install the -cu12 wheels). See the ONNX Runtime CUDA requirements.

The GPU path is tested here with mocked sessions and on a GPU-less host with a real onnxruntime-gpu (it raises / warns as documented); numerical GPU-vs-CPU agreement has not been measured yet in this package (TF32 is disabled by default, see cuda_tf32).

class brainmaze_eeg_models.runtime.OnnxModel(path: str | PathLike | bytes, *, device: str = 'auto', threads: int | None = None, batch_size: int = 32, cuda_device_id: int = 0, cuda_tf32: bool = False, output_batch_axes: Mapping[str, int] | None = None, sha256: str | None = None, name: str | None = None)#

One ONNX model on one device.

Parameters:

  • path (str, os.PathLike or bytes) – Path to the .onnx file, or the serialized model itself as bytes (as onnxruntime.InferenceSession accepts; then name is required, sha256 is checked on the bytes and path is None).

  • device ({'auto', 'cpu', 'cuda'}) – 'auto': CUDA when the CUDA execution provider is installed and initialises, else CPU. A RuntimeWarning is issued when the CPU is used although a GPU setup was evidently intended: onnxruntime-gpu is installed but CUDA cannot be initialised, or both onnxruntime and onnxruntime-gpu are installed and the CPU build is the active module. 'cpu': always CPU. 'cuda': CUDA or RuntimeError (never a silent CPU fallback).

  • threads (int, optional) – CPU threads per inference call (intra_op_num_threads). None: ONNX Runtime default (one per physical core). Also used for the CPU parts of a CUDA session.

  • batch_size (int) – Default number of items per inference call in run().

  • cuda_device_id (int) – Index of the GPU used with device='cuda' / 'auto' (default 0).

  • cuda_tf32 (bool) – Allow TF32 matrix maths on the GPU. Default False: ONNX Runtime’s CUDA provider enables TF32 by default, which changes results at the ~1e-3 level (not the ~1e-6 of float32 rounding) on Ampere and newer GPUs; off, GPU and CPU outputs are expected to agree to float32 rounding (not yet measured on a GPU for these models). Must be a bool.

  • output_batch_axes (mapping name -> int, optional) – Batch axis of each output, if not 0 (e.g. {'probs': 1} for a (time, batch, ...) output). run() checks on every chunk that this axis has the chunk’s length and concatenates the chunks along it; a wrong axis raises instead of silently mixing items.

  • sha256 (str, optional) – Expected SHA-256 hex digest of the file; a mismatch raises RuntimeError.

  • name (str, optional) – Name used in messages.

property providers: list[str]#

Execution providers of the session, in priority order.

run(inputs: Mapping[str, ndarray], batch_size: int | None = None) → dict[str, ndarray]#

Run the model on a batch, in chunks along the first axis.

Parameters:

  • inputs (mapping name -> array) – One array per model input (input_names), all with the same length of the first (batch) axis, at least 1. Cast to the model’s dtype chunk by chunk; must be finite (checked for the whole batch before the first inference call).

  • batch_size (int, optional) – Items per inference call; default batch_size. Results do not depend on it beyond float32 rounding (a few ulp; the kernels may block differently).

Returns:

One array per model output (output_names), concatenated along its batch axis (0 unless set with output_batch_axes).

Return type:

dict name -> np.ndarray

Raises:

RuntimeError – If an output’s batch axis does not have the length of the chunk (wrong output_batch_axes), so that items are never silently mixed up.

brainmaze_eeg_models.runtime.available_providers() → list[str]#

Execution providers compiled into the installed onnxruntime (not necessarily usable).

brainmaze_eeg_models.runtime.device_report() → str#

Human-readable summary of the installed onnxruntime distributions and providers.

brainmaze_eeg_models.runtime.check_gpu(device_id: int = 0) → str#

Check that CUDA really works: create a CUDA session and run a tiny inference on the GPU.

Unlike onnxruntime.get_available_providers() (which lists CUDAExecutionProvider whenever the GPU build is installed, even if the CUDA libraries are missing and every session would fall back to the CPU), this runs a small model with device='cuda' on GPU device_id and compares its output with the exact result.

Returns:

A one-line confirmation (GPU index, session providers, onnxruntime version).

Return type:

str

Raises:

RuntimeError – If CUDA cannot be used, with the reason and what to install (as device='cuda').