Skip to content

feat(inference): prefer ONNX Runtime with PyTorch fallback#126

Merged
Goldokpa merged 1 commit into
Climate-Vision:mainfrom
UTKARSH698:feat/onnx-inference-path
Jul 23, 2026
Merged

feat(inference): prefer ONNX Runtime with PyTorch fallback#126
Goldokpa merged 1 commit into
Climate-Vision:mainfrom
UTKARSH698:feat/onnx-inference-path

Conversation

@UTKARSH698

Copy link
Copy Markdown
Contributor

Summary

Completes #12 by making ONNX Runtime the preferred serving backend, with a clean fallback chain, per-request latency observability, and a batch export path.

Inference (inference/pipeline.py)

  • _load_model now prefers an exported .onnx model when one is present and loadable, then falls back to the PyTorch .pth checkpoint, and finally to untrained demo weights. Previously ONNX was only tried when no checkpoint existed, so exported models were never served.
  • Every request logs and returns its backend and forward-pass latency: the inference block of the /api/predict response gains backend (onnx/pytorch) and inference_ms.
  • New benchmark_inference(analysis_type) helper times both backends side by side and reports the speedup.

Export (scripts/export_model.py)

  • New --all flag exports every enabled analysis type to models/unet_<type>.onnx, reading each type's checkpoint path and num_classes from config.yaml (so 3-class ice/flood models export correctly, not just 2-class deforestation).
  • validate_onnx now asserts the (N, n_classes, H, W) output signature, catching shape regressions at export time.
  • Single-checkpoint mode is refactored onto the same export_checkpoint helper; --checkpoint is no longer required when --all is used.

Tensor signatures verified

UNet (N, C, 256, 256)(N, n_classes, 256, 256), with dynamic batch/H/W axes preserved.

Test plan

  • tests/test_onnx_inference.py (new): ONNX preferred over .pth; fallback when no ONNX; fallback when ONNX unloadable (LFS pointer); backend/inference_ms surfaced in the response; benchmark_inference shape; end-to-end export→load→serve roundtrip; --all discovery covers enabled types.
  • Existing tests/test_pipeline.py still green; full suite shows no regressions (pre-existing flood/SAR failures are unrelated — they need the optional scikit-image dep).

Closes #12

_load_model now serves an exported .onnx model when one is present and
loadable, falling back to the PyTorch .pth checkpoint and finally to
untrained demo weights. Each request logs and returns its backend and
forward-pass latency (backend, inference_ms), and benchmark_inference()
times both backends side by side.

export_model.py gains --all to export every enabled analysis type to
models/unet_<type>.onnx using each type's num_classes from config.yaml,
and validate_onnx now asserts the (N, n_classes, H, W) output signature.

@Goldokpa Goldokpa left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for this, @UTKARSH698 — nice work, especially catching that the old logic only tried ONNX when no .pth checkpoint existed, so exported models were silently never served. The fixed priority order (ONNX -> PyTorch -> untrained demo) closes that gap cleanly.

A few things I liked: the export_checkpoint() refactor is solid, since single-checkpoint and --all batch export now share one code path instead of duplicating logic; the output-shape assertion in validate_onnx is a good guardrail that catches a (N, n_classes, H, W) mismatch at export time rather than at serve time; surfacing backend and inference_ms in the /api/predict response is a nice observability win, and benchmark_inference() makes the ONNX vs PyTorch tradeoff easy to verify empirically; and test coverage looks thorough, covering preference order, fallback on missing or unloadable ONNX, and the export-load-serve roundtrip.

Nothing blocking. Approving.

@Goldokpa
Goldokpa merged commit 40058de into Climate-Vision:main Jul 23, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Good First Issue] Add ONNX Runtime inference path with PyTorch fallback

2 participants