Feature caching for the interactive annotate/predict loop - #73
Open
hinderling wants to merge 9 commits into
Open
Feature caching for the interactive annotate/predict loop#73hinderling wants to merge 9 commits into
hinderling wants to merge 9 commits into
Conversation
Generic, FE-agnostic cache for the interactive annotate→predict loop: stores opaque per-image payloads keyed by content hash + FE signature, bounded by a RAM budget with live headroom (never OOMs — degrades to recomputation), LRU eviction, and an optional disk tier that RAM-evicted or oversized payloads spill to (per-entry opt-out via spill_ok for payloads cheaper to recompute than to pickle).
Split extract_features_pyramid into _pyramid_native (expensive per-scale native features, device-independent CPU numpy — the cacheable payload) and _pyramid_reconstruct (cheap rescale+concat), keeping the combined output bit-identical. Add the cache protocol hooks on the base class (supports_feature_cache / cacheable_repr / features_from_cacheable / cacheable_nbytes / cache_spill_to_disk / cache_extra_state) and the FE overrides: Hookmodel keeps its huge CNN payloads out of the disk tier, ComboFeatures opts out (it bypasses the pyramid split), and JAFAR/gaussian report their instance state (jafar_scalings, sigma) for the cache key. Includes two micro-fixes the split exposed: features_per_layer now defaults to None on the base class (fe_use_min_features degrades with a warning instead of AttributeError), and supported_devices no longer returns a nested list for the CPU entry.
enable_feature_cache() attaches a FeatureCache; _get_features then routes per-image pyramid extraction through _extract_pyramid_cached, which keys the FE's native payload by (content hash of the prepared image, FE signature). The signature covers the train-reset param set plus image_downsample, patch size, and FE instance state (cache_extra_state), so any setting that changes feature values misses instead of serving stale features. With the cache disabled (default) behaviour is unchanged; enabled, outputs are bit-identical (the pyramid split is exact). cache_only=True threads a peek through _predict/_predict_image/ _get_features: return the prediction only if all needed features are already cached, else None — lets a stack prediction serve cached slices first, before a sequential scan evicts them (used by the widget).
GUI wiring for the cache: enable checkbox + RAM/disk budget settings, live size/hit-rate label, (re)creating the cache when the model changes, and clearing it only when the last user image layer is removed (plugin- owned probability/feature layers — including renamed backups — never wipe it; content-addressing handles staleness otherwise). Predict-all now runs cache-first: a cheap peek pass serves all slices whose features are already cached before the compute pass runs — a plain sequential scan over a stack larger than the cache would evict exactly the slices it is about to need (LRU thrash). Extended the cache tests with the model-level protocol (disk routing, spill opt-out, FE-state keys, bit-identical cached predictions).
The widget already normalizes the full stack (image_stack_norm) before training, and prediction passes skip_norm=True on the same pre-normalized data — but training passed skip_norm=False, so the model ran a second normalization pass. That pass never changed values meaningfully (verified per mode: the imagenet guard returns out-of-[0,1] data unchanged with a spurious warning; default z-scoring and the percentile stretch are both idempotent to float32 noise), but it wasted a full-stack pass, emitted a bogus warning for imagenet FEs, and left train/predict arrays differing at the float-noise level — which breaks feature-cache sharing, since cache keys are content hashes and need byte-identical inputs. Passing skip_norm=True makes train and predict prepare exactly the same arrays. The multi-file training path still passes skip_norm=False — it feeds raw, not-yet-normalized images.
hinderling
marked this pull request as ready for review
July 13, 2026 15:56
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #73 +/- ##
==========================================
- Coverage 63.20% 61.15% -2.06%
==========================================
Files 32 34 +2
Lines 6515 7162 +647
==========================================
+ Hits 4118 4380 +262
- Misses 2397 2782 +385 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
A worker thread can now extract/cache while the GUI thread clears the cache or changes limits (needed for the threaded-cancellation branch, and already correct for API users sharing a model across threads).
… tensor-aware pad_to_shape skimage.transform.resize computes a full resize even when the target shape equals the input shape — for the common patch_size=1 / image_downsample=1 prediction path that was ~17 ms/slice of pure waste (rescale_features already had this guard). The class-labels guard keeps the uint8 contract. pad_to_shape gains a torch branch (same symmetric convention) so it works for on-device feature tensors.
The pyramid split cast NN features to CPU numpy inside _pyramid_native, which silently moved the whole multi-scale rescale from torch-on-device onto CPU skimage (rescale_features dispatches on array type): ~2-4x slower stack prediction than main on GPU, cache on or off. - _pyramid_native keeps features in their native form (torch tensors on the extraction device for NN FEs); _pyramid_reconstruct rescales and concatenates on-device with a single host transfer at the end. - The cache payload is device-independent CPU numpy (_native_to_payload) and records was_torch; features_from_cacheable lifts such payloads back onto the device, so cache HITS use the same torch backend as fresh extractions: a hit is now faster than a recompute (it saves the forward pass) instead of ~5x slower, and hit/miss results stay identical. - On a miss, _extract_pyramid_cached reconstructs from the on-device native form via cacheable_repr_and_features (one extraction pass) and stores only the numpy payload, so the cache-enabled path is as fast as the cache-off path. - Numpy-native FEs (gaussian etc.) are untouched: their payloads stay numpy and reconstruct via skimage on hit and miss alike. - Caching surface trimmed while at it: cacheable_repr and cacheable_nbytes dropped (cacheable_repr_and_features is the single extension point; the cache sizes payloads itself), invalidate(predicate) folded into clear() (content-addressed entries never go stale), disk_dir ctor param dropped, cache_extra_state takes no Param, the widget's three cache-settings handlers collapse into _apply_feature_cache, and _predict_all's two loops merge into one cache-first pair. Benchmarks (VGG16, scalings [1,2,4], MPS): main 34 ms/slice; before this fix 125 (cache on or off); after: 18.3 off / 19.8 miss / 15.4 hit — bit-identical across all paths and devices. Tests: hit == miss == cache-off equality for an NN FE (vgg16) and a numpy FE, payload-form assertions, plus a threaded cache hammer.
Contributor
Author
|
Pushed three commits addressing a performance regression found while benchmarking this branch on GPU:
Benchmarks (VGG16, scalings [1,2,4], MPS): main 34 ms/slice → 18.3 cache-off / 19.8 miss / 15.4 hit (full table in the last commit message). Verified against the #69 tiling-consistency suite (17/17) plus cache-on tiled-vs-untiled and non-patch-multiple DINOv2 probes — all pixel-identical. |
…t site The docstring already documents the device contract; add an inline marker at the actual fallback (hoisted into lift_device) so the caveat is visible where the hit is reconstructed. No behavior change.
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
2 participants
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
Opt-in feature caching for the interactive annotate → predict loop: when the same image (or z-slice / movie frame) is segmented repeatedly while refining scribbles, its extracted features are reused instead of recomputed. This is by far the biggest win for ViT-family extractors (DINOv2/v3, JAFAR), where feature extraction dominates the loop — a re-predict on a cached slice is near-instant.
Cached output is bit-identical to uncached extraction: the cache stores the native per-scale features (the expensive half of the pyramid), and the cheap rescale/concat half runs on every request. Entries are content-addressed (blake2b of the prepared image + a signature of every feature-relevant setting), so the cache is self-invalidating — a changed image or setting simply misses; nothing needs manual invalidation.
Design (one commit per layer, reviewable top to bottom)
feature_cache.py— a generic, FE-agnostic LRU cache bounded by a RAM budget with live headroom (it degrades to recomputation rather than OOM), plus an optional disk tier: RAM-evicted or RAM-oversized payloads spill to disk (with a per-entry opt-out for payloads cheaper to recompute than to pickle).extract_features_pyramidis split into_pyramid_native(cacheable, device-independent) and_pyramid_reconstruct(cheap), with protocol hooks on the base class. Per-FE specifics: CNN Hookmodels keep their huge per-pixel payloads out of the disk tier, ComboFeatures opts out, and JAFAR/gaussian report instance state (jafar_scalings,sigma) for the cache key so changing it can never serve stale features.ConvpaintModel—enable_feature_cache()plus the key derivation and acache_only"peek" mode threaded through the predict path.Tests
test_feature_cache.py: 20 tests covering the cache module (LRU order, budgets, disk spillover, invalidation) and the model-level protocol (disk routing of oversized payloads, spill opt-out, FE-state keys, peek semantics, and bit-identical cached predictions). Full suite passes (the two cellpose tests fail onmaintoo in my environment — broken local cellpose install, unrelated).Extracted from the larger performance work in #72 to be reviewable on its own; a follow-up PR with widget layout polish builds on this one.