feat(inference): prefer ONNX Runtime with PyTorch fallback#126
Conversation
_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
left a comment
There was a problem hiding this comment.
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.
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_modelnow prefers an exported.onnxmodel when one is present and loadable, then falls back to the PyTorch.pthcheckpoint, and finally to untrained demo weights. Previously ONNX was only tried when no checkpoint existed, so exported models were never served.inferenceblock of the/api/predictresponse gainsbackend(onnx/pytorch) andinference_ms.benchmark_inference(analysis_type)helper times both backends side by side and reports the speedup.Export (
scripts/export_model.py)--allflag exports every enabled analysis type tomodels/unet_<type>.onnx, reading each type's checkpoint path andnum_classesfromconfig.yaml(so 3-class ice/flood models export correctly, not just 2-class deforestation).validate_onnxnow asserts the(N, n_classes, H, W)output signature, catching shape regressions at export time.export_checkpointhelper;--checkpointis no longer required when--allis 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_mssurfaced in the response;benchmark_inferenceshape; end-to-end export→load→serve roundtrip;--alldiscovery covers enabled types.tests/test_pipeline.pystill green; full suite shows no regressions (pre-existing flood/SAR failures are unrelated — they need the optionalscikit-imagedep).Closes #12