feat(inference): prefer ONNX Runtime with PyTorch fallback#126
Open
UTKARSH698 wants to merge 1 commit into
Open
feat(inference): prefer ONNX Runtime with PyTorch fallback#126UTKARSH698 wants to merge 1 commit into
UTKARSH698 wants to merge 1 commit into
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.
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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