perf(py): normalise the patch blend per share, not by an accumulated weight - #69
Merged
Conversation
…weight Each patch is now blended in already carrying its fraction of the blend weight, w / sum_k w, taken from a per-axis factorisation. The shares sum to one per voxel by construction, so the spatial-sized weight buffer and the final division pass over every channel both disappear. The factorisation is exact, not an approximation: the patch grid is a full per-axis product and the window is separable by construction, so sum_p prod_d w_d == prod_d sum_k w_d. The total weight is therefore one vector per axis -- 13 KB where a 320-row window of a 1331x1775 volume held 2.8 GB. Measured on a whole-body CT (295x259x219, patch [96,128,160]), the two regimes that bracket the channel count: 117-class segmentation head, fp16 42.7s -> 39.7s RSS 6.07 -> 6.02 GiB single-channel synthesis, fp32 2.2s -> 2.0s single-channel, streamed 3.0s -> 2.4s The memory saved is the weight buffer, 1/(C+1) of the accumulator: negligible on a 117-class head, 25% on a 3-component field, 50% on a single channel (0.237 -> 0.090 GiB of VRAM on a 128x384x384 case, the final division going with it). It is also more accurate in low precision. The weight no longer needs a floor: the share is a ratio of comparable quantities and stays in [0, 1], where the raw product underflows fp16 (a Gaussian corner is ~1e-8 against a 6e-5 smallest normal) and had to be clamped on both sides of the division. Against a float64 reference the assembled volume is 519x closer on that corner, and the fp16 whole-volume std lands 53x nearer its float32 value. float32 output is unchanged to 3.6e-07, and Mean and the unweighted path are bit-identical. Two tests pin what the change rests on: that the accumulated weight factorises into one vector per axis, and that a float16 blend returns the volume it was cut from, borders included -- both fail on the accumulated form.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughPatch blending now derives separable per-axis geometry and applies blend shares during accumulation. Full-volume weight buffers and final normalization were removed from regular and streaming accumulators, with tests covering factorization and float16 reconstruction. ChangesPatch blending refactor
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant PatchInput
participant Accumulator
participant ResultBuffer
PatchInput->>Accumulator: add cropped patch
Accumulator->>Accumulator: compute blend share
Accumulator->>ResultBuffer: add share-weighted patch
Accumulator->>ResultBuffer: assemble without weight normalization
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
konfai/data/patching.py (1)
273-289: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUpdate the blending docs and drop the dead weight path
weight()and theMean/Cosinus/Gaussiancomments still describe the oldweight_sumnormalization; rewrite them to match the current share-based assembly.PathCombine.weight()/__call__()also have no in-repo callers, so remove the dead path and_data_per_devicecache if nothing external depends on them.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@konfai/data/patching.py` around lines 273 - 289, Update the blending documentation and comments in weight() and the Mean, Cosinus, and Gaussian implementations to describe the current share-based assembly rather than weight_sum normalization. Remove PathCombine.weight(), PathCombine.__call__(), and the _data_per_device cache, provided repository-wide usage confirms no callers or dependencies remain.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@konfai/data/patching.py`:
- Around line 273-289: Update the blending documentation and comments in
weight() and the Mean, Cosinus, and Gaussian implementations to describe the
current share-based assembly rather than weight_sum normalization. Remove
PathCombine.weight(), PathCombine.__call__(), and the _data_per_device cache,
provided repository-wide usage confirms no callers or dependencies remain.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: e22f4de9-d919-4755-a78a-1181d6af9770
📒 Files selected for processing (2)
konfai/data/patching.pytests/unit/test_patching.py
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.
What
Each patch is blended in already carrying its fraction of the blend weight,
w / sum_k w, taken froma per-axis factorisation. The shares sum to one per voxel by construction, so the spatial-sized
weight buffer and the final division pass over every channel both disappear.
Why it is exact
The patch grid is a full per-axis product and the window is separable by construction, so summing it
over the grid factorises:
The total weight is therefore one vector per axis — 13 KB where a 320-row window of a 1331×1775
volume held 2.8 GB.
Measured
Whole-body CT, 295×259×219, patch
[96, 128, 160]— the two regimes that bracket the channel count:The memory saved is the weight buffer,
1/(C+1)of the accumulator, so it depends entirely on thechannel count: negligible (0.8%) on a 117-class head, 25% on a 3-component field, 50% on a single
channel — 0.237 → 0.090 GiB of VRAM on a 128×384×384 case, the final division going with it.
One regime is slower: a single-channel streamed blend on a small volume (5.7 → 12.4 ms), where three
broadcast multiplies cost more than one multiply by a precomputed 3-D window and there is almost no
final division to save. It inverts at application scale (1.26× faster above).
Accuracy
The weight no longer needs a floor. The share is a ratio of comparable quantities and stays in
[0, 1], where the raw product underflows fp16 — a Gaussian corner is ~1e-8 against a 6e-5 smallestnormal — and had to be clamped on both sides of the division to stay recoverable.
Meanand the unweighted path bit-identical.the 117-class case lands 53× nearer its float32 value (0.243624 → 0.244509, reference 0.244526).
Tests
Two tests pin what the change rests on, and both fail on the accumulated form:
Existing suite green (
ruff,ruff format, mypy, bandit,tests/unit).Summary by CodeRabbit
Performance
Bug Fixes
Tests