ONNX export: optional early stopping of non-covisible pairs - #7

Stacked on #6 (fused GELU) and #5. This PR adds optional early stopping of non-covisible pairs to the ONNX matcher, as FlexGlue.match(..., early_stop_decay=...) does in PyTorch.

What

• python onnx/export.py --early_stop_layers 2,4,6,8 exports flexglue_matcher_early_stop{,_fp16}.onnx.
• After each listed layer k, the pair is rejected without running the remaining transformer layers and the readout if its covisibility score is below min_covisibility * (1 - early_stop_decay) ** (9 - k).
  • In that case, matches0 = -1, mscores0 = 0, and covisibility_score is the score after layer k.
• early_stop_decay is a new scalar input. early_stop_decay = 1 (or min_covisibility = 0) never stops early, and then gives the same results as the regular matcher.
• onnx/match_pair.py gets --early_stop (selects the early-stop models) and --early_stop_decay, and passes early_stop_decay to models that have it. The README documents the feature.

How

• torch.cond doesn't work here. In torch 2.6 it needs dynamo to trace the branches, which fails on Tensor.unflatten and on the custom ops' autograd. With these fixed, tracing 8 nested conditionals ran for more than 15 minutes without finishing.
• Instead, the If nodes are built on the exported ONNX graph (add_early_stopping):
  1. The model is exported with additional per-layer covisibility-score outputs (LayerScoresMatcherONNX).
  2. The nodes are split into stages: the nodes needed for the score after the first checked layer, those needed in addition for the next one, and so on, and finally the readout.
  3. Each stage after a check becomes the else-branch of an If whose then-branch outputs the rejection.
  4. Each initializer moves into the subgraph of its first user.
• The export takes about 20 s.

Does the If actually skip the computation? Yes.

fp16, N = 2048, A100, ONNX Runtime 1.23. Wall time per pair, device I/O:

• Each transformer layer costs about 3.3 ms, and the work after a stop is skipped entirely.
• Cost when a pair doesn't stop:
  • GPU time: +0.3 ms per pair with 4 checks.
  • Wall time: about +0.2 ms per check, because ONNX Runtime needs each condition on the host and synchronizes.
  • The fused GELU (ONNX export: fused GELU operator in the matcher #6) partly compensates for it. In COLMAP, pairs that never stop take 54.5 vs. 53.8 ms/pair at 2k (+1%).

How much do we gain? It depends on the share of non-covisible pairs.

Per-layer covisibility scores of all exhaustive pairs, checks after layers 2/4/6/8, early_stop_decay=0.5:

• Lower decays stop more often but reject true pairs:
  • decay 0.3: −18% to −23% on ETH3D, losing 1.8% of the kept pairs;
  • decay 0.2: about −30%, losing 4–5%;
  • decay 0.5 is the safe choice.
• End to end in COLMAP (3 ETH3D scenes, 2132 pairs; total matching time / pairs, including image decoding):
  • 2k keypoints: −5.7% (min_covisibility 0.5) and −7.3% (0.8).
  • 8k keypoints: −9.2% and −12.9%, since the readout skipped by a stop is more expensive.

Parity

• onnx/check.py asserts the early-stop branch:
  • It always compares with early_stop_decay=1, which never stops.
  • If the graph checks the covisibility after every layer but the last, as FlexGlue does (--early_stop_layers 1,...,8), it also compares matches0, the scores and covisibility_score with FlexGlue.match(..., early_stop_decay=d) for d ∈ {0, 0.2, 0.5}. On the asset pairs, the stops happen after different layers.
  • Results: identical on CPU; passes on CUDA in fp32 and fp16 at 512, 2048 and 4096 keypoints.
  • For other layer sets, e.g. the released 2/4/6/8 models, it prints the outputs with early stopping.
• README: release notes on the v0.1 matchers. The early-stop models have the fused GELU and checks after layers 2/4/6/8. The base matchers were exported before ONNX export: fused GELU operator in the matcher #6, without the fused GELU.
• CI (Add ONNX export and an ONNX Runtime example #5): also exports the shipped 2/4/6/8 early-stop configuration (export only, ~30 s), so it keeps building.
• CI (Add ONNX export and an ONNX Runtime example #5): also exports the float32 early-stopping matcher with layers 1–8 and checks it on CPU. This adds about 8 minutes (export 0.5 min, check 7.6 min), for a total of about 14.5 minutes.
• With decay 1, the outputs match the regular matcher.
• With decay 0.5, the covisible pair keeps all its matches and the look-alike pair stops after layer 6.

The models are flexglue_matcher_early_stop.onnx and flexglue_matcher_early_stop_fp16.onnx (sha256 c86714cd… / 01286e69…), to be uploaded to the v0.1 release for COLMAP.