Skip to content

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

Open
UTKARSH698 wants to merge 1 commit into
Climate-Vision:mainfrom
UTKARSH698:feat/onnx-inference-path
Open

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

Conversation

@UTKARSH698

Copy link
Copy Markdown

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.
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

1 participant