Skip to content

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

Closed

Description

@Oshgig

Overview

scripts/export_model.py exports PyTorch models to ONNX, but the API still loads .pth checkpoints directly. ONNX Runtime would provide faster inference and lower memory usage.

Scope

  • Update _load_model() in src/climatevision/inference/pipeline.py to prefer .onnx files when present
  • Fall back to PyTorch .pth if ONNX file is missing
  • Verify ONNX input/output tensor signatures match per analysis type:
    • UNet: (N, n_channels, 256, 256) → (N, n_classes, 256, 256)
    • Siamese: (N, C, 256, 256) × 2 → (N, 2, 256, 256)
  • Add latency benchmark logging (compare PyTorch vs ONNX inference time)
  • Update scripts/export_model.py to export all analysis-specific models

Acceptance Criteria

  • When .onnx checkpoint exists, API loads it via ONNX Runtime
  • Latency difference is logged per request
  • Fallback to PyTorch works seamlessly when ONNX export is absent
  • All 3 analysis types supported

Resources

  • src/climatevision/inference/pipeline.py — _load_model() function
  • scripts/export_model.py — existing ONNX export script
  • references/model-architectures.md — tensor shape specs

Difficulty: Intermediate
Labels: help wanted, backend, mlops, performance

Activity

  1. added
    help wantedExtra attention is needed
    and removed
    help wantedExtra attention is needed
    on Apr 19, 2026
  2. changed the title [-][Help Wanted] Add ONNX Runtime inference path with PyTorch fallback[/-] [+][Good First Issue] Add ONNX Runtime inference path with PyTorch fallback[/+] on Apr 25, 2026
  3. UTKARSH698 commented on Jul 22, 2026

    @UTKARSH698
    Contributor

    Opened #126 for this.

    It makes ONNX Runtime the preferred backend when an exported .onnx model is present (falling back to the .pth checkpoint, then untrained demo weights), adds per-request latency logging — the /api/predict response now carries backend and inference_ms — and a benchmark_inference() helper that times both backends side by side. scripts/export_model.py gains --all to export every enabled analysis type to models/unet_<type>.onnx (reading each type'''s num_classes from config.yaml), and validate_onnx now asserts the (N, n_classes, H, W) output signature. Covered by a new tests/test_onnx_inference.py.

    Happy to adjust the load-priority default or the response fields if you'''d prefer them shaped differently.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions