feat: ONNX export and torch-free runtime for the tube classifier - #76
Merged
Merged
Conversation
…backend Move predict() into TubePipelineModel, keep torch only in the BboxTubeTemporalModel scoring and detection overrides. Patches are built in numpy, the trigger search takes a numpy score callable, and StageTimer only imports torch for accelerator syncs.
temporal-export-onnx derives model_onnx.zip (fixed-shape classifier, config, calibrator, source sha256 and I/O contract) from a model.zip and checks torch/onnxruntime parity on several mask patterns. OnnxTemporalModel runs the shared pipeline on supplied detections without importing torch.
The base package is now numpy/pillow/pydantic/pyyaml only. Training, packaging and the torch backend use [torch]; the onnxruntime backend uses [onnx]. Dependent packages pin temporal-model-core[torch]; the core dev group installs both extras for the test suite.
publish --onnx-file uploads the exported archive under the same HF tag, fetch --onnx downloads it. Add make export-onnx / fetch-model-onnx and document the [torch] / [onnx] extras and the ONNX runtime usage.
torch 2.12's dynamo exporter needs onnxscript >= 0.7 (onnx-ir with schemas); the export's parity check needs onnxruntime, so the [torch] extra carries it. numpy < 2 has no 3.13 wheels, and newer onnxscript made the universal lock fail on that split, so say so in requires-python.
Stamping the version rewrites model.zip, so the SHA-256 recorded at export time no longer identified the released archive.
…model.zip Compare the manifest's recorded source hash with the input model.zip before re-stamping it from the version-stamped copy.
…e runbook torch.onnx.export(dynamo=True, external_data=False) needs 2.5. The runbook chained cd commands that ended up in core/api.
- export_classifier/verify_export work on deep copies: the caller's module is never moved to CPU or de-fused, and parity is now measured against the as-served (fused-attention) torch path, not the export-friendly variant. This also fixes the order-dependent module-scoped fixture in test_onnx.py. - Export with a dynamic batch axis (batch-2 example input so the tracer cannot fold the size-1 batch dim): OnnxTemporalModel._score is one session.run for all tubes instead of a per-tube Python loop, and verify_export adds a batched parity pass. - build_onnx_package refuses an uncalibrated model.zip by default (--allow-uncalibrated to override), instead of shipping an artifact that only fails on the edge device after an immutable tag is cut. - Re-export UncalibratedModelError from onnx_model so [onnx]-only installs have a torch-free import path for it. - StageTimer validates/normalizes the device via torch.device when torch is installed (a typo like "gpu" raises instead of silently disabling sync). - load_model_package(with_detector=False) no longer extracts YOLO weights. - release.py: reuse core's ONNX_MODEL_FILENAME, stream SHA-256 with hashlib.file_digest, stamp the staged ONNX manifest in one rewrite. - torch floor to >= 2.6 (external_data kwarg); torchvision declared in core's dev group for the parity test; providers=None vs [] fixed. - docs: single-publish release runbook (the tag is immutable, a second publish cannot add the ONNX archive), drop stale score_tubes references, add docs/specs/2026-09-09-onnx-export-design.md referenced from the new module docstrings.
MateoLostanlen
force-pushed
the
feat/onnx-export
branch
from
September 10, 2026 05:52
ef76860 to
9430a96
Compare
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
1 participant
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.
model_onnx.zip, a torch-free runtime artifact published next tomodel.zipunder the same HF tag.predict()moves into a backend-agnosticTubePipelineModel;BboxTubeTemporalModelkeeps torch and YOLO, the newOnnxTemporalModelruns on onnxruntime with supplied detections.temporal-export-onnxderives the archive with fixed shapes, checks torch/ONNX parity on several mask patterns and records the source hash and I/O contract in its manifest.release publish --onnx-fileandfetch --onnxship it.[torch]and[onnx]extras; the base package is numpy, pillow, pydantic and pyyaml. Verified on a Pi 5: same logits as torch to 1e-5, 161 MB install.