onnx_docker_compatibility

Compatibility test bench for the ONNX models used in boom (ZTF alert classification).

The goal is to find the correct CUDA / cuDNN versions needed to run these models on GPU inside Docker, before wiring GPU inference into the main pipeline.

Models

All models live in data/models/ and are copied from the boom repository.

File	Task	Inputs	Output
acai_h/n/v/o/b.d1_dnn_20201130.onnx	Real-bogus / source classification (5 variants)	features (B, 25) + triplets (B, 63, 63, 3)	score (B,)
btsbot-v1.0.1.onnx	Bright transient classification	metadata (B, 25) + triplet (B, 63, 63, 3)	fc_out (B,)

B = batch size. All tensors are f32.

What the binary does

1. Loads all 6 model sessions (CPU or GPU execution provider).
2. Creates dummy inputs filled with constant values — no real alert data needed.
3. Runs inference at batch sizes 1, 8, and 32.
4. Prints session load time, inference time, and the first score from each model.

This lets you confirm that:

• The models load correctly under a given ORT / CUDA / cuDNN combination.
• GPU inference actually runs on the device (check the ORT log for CUDA EP messages).
• Batch throughput scales as expected before integrating into boom.

Version matrix

Component	Version
ort crate	2.0.0-rc.10
ONNX Runtime	1.22.0
CUDA	12.8
cuDNN	9.20

Local build (CPU only)

cargo run

No CUDA required. The GPU code path compiles to a no-op warning when the gpu feature is absent.

Local build (GPU)

Requires ORT 1.22.0 with CUDA support installed and ORT_LIB_LOCATION set.

ORT_LIB_LOCATION=/path/to/onnxruntime/lib cargo run --features gpu

Docker

# CPU build (default)
docker build -t onnx-compat .
docker run onnx-compat

# GPU build — downloads CUDA 12.8 + cuDNN 9.20 from NVIDIA redist
docker build --build-arg USE_GPU=true -t onnx-compat-gpu .
docker run --gpus all onnx-compat-gpu

The builder stage is rust:slim-bookworm. CUDA/cuDNN are installed from NVIDIA's official redist archives only when USE_GPU=true, keeping the CPU image small. The ORT shared library bundled by the ort crate (libonnxruntime.so) is copied into the final debian:bookworm-slim image alongside the binary.

Output example

=== CPU batch (size=1) ===
  Session load : 29.6ms
  Inference    : 5.2ms
  acai_h[0]    = 0.000000
  acai_n[0]    = 0.000002
  acai_v[0]    = 1.000000
  acai_o[0]    = 0.000000
  acai_b[0]    = 0.802032
  btsbot[0]    = 0.000000

=== GPU batch (size=1) ===
  Session load : 312.4ms
  Inference    : 1.1ms
  acai_h[0]    = 0.000000
  ...

Related

• boom — the main pipeline this feeds into
• ort — Rust bindings for ONNX Runtime
• ONNX Runtime CUDA EP docs