Optimize FFN segmentation loading and cleanup for empty and dense subvolumes. - #129
Open
copybara-service[bot] wants to merge 1 commit into
Open
Optimize FFN segmentation loading and cleanup for empty and dense subvolumes.#129copybara-service[bot] wants to merge 1 commit into
copybara-service[bot] wants to merge 1 commit into
Conversation
…volumes. In large-scale sparse pipelines (like whole-brain agglomeration), many subvolumes are empty (contain only background). Running the full connected components and size filtering (clean_up) on these empty arrays is wasteful. This CL introduces three optimizations with safe corner-case handling: 1. **Early exit in `load_segmentation`**: If the loaded segmentation from NPZ is empty, we return early. This avoids a 512MB allocation/cast (`astype(np.uint64)`) and avoids calling `clean_up` entirely. 2. **Empty Block Fast Path in `clean_up` / `clean_up_and_count`**: Checks if the array is empty using `np.any` (takes ~50ms instead of ~1.7s on 400x400x400 empty array). If empty, returns early with properly typed mappings (handling zero-sized arrays as well). 3. **Linear-time Size Filtering (`clear_dust`)**: For non-empty integer blocks, if the max segment ID is small (< 10M, which is typical for local subvolume IDs before global relabeling), replaces `np.unique` (which sorts 64M elements) with `np.bincount` and lookup-table indexing matching `data.dtype` to avoid memory expansion. This yields up to **17x speedup for dense blocks** (from ~3.0s to ~0.17s for 90% density). Properly supports all integer dtypes (including uint8/uint16 without overflow) and signed arrays (including negative segment IDs without positive max), with graceful fallback to `np.unique` for non-integer arrays. PiperOrigin-RevId: 947892396
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
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.
Optimize FFN segmentation loading and cleanup for empty and dense subvolumes.
In large-scale sparse pipelines (like whole-brain agglomeration), many subvolumes are empty (contain only background). Running the full connected components and size filtering (clean_up) on these empty arrays is wasteful.
This CL introduces three optimizations with safe corner-case handling:
load_segmentation: If the loaded segmentation from NPZ is empty, we return early. This avoids a 512MB allocation/cast (astype(np.uint64)) and avoids callingclean_upentirely.clean_up/clean_up_and_count: Checks if the array is empty usingnp.any(takes ~50ms instead of ~1.7s on 400x400x400 empty array). If empty, returns early with properly typed mappings (handling zero-sized arrays as well).clear_dust): For non-empty integer blocks, if the max segment ID is small (< 10M, which is typical for local subvolume IDs before global relabeling), replacesnp.unique(which sorts 64M elements) withnp.bincountand lookup-table indexing matchingdata.dtypeto avoid memory expansion. This yields up to 17x speedup for dense blocks (from ~3.0s to ~0.17s for 90% density). Properly supports all integer dtypes (including uint8/uint16 without overflow) and signed arrays (including negative segment IDs without positive max), with graceful fallback tonp.uniquefor non-integer arrays.