Repository navigation
Conversation
Keep low-resolution logits only for mask-to-mask refinement, their last consumer. Remove them after refinement so later filters, concatenation and CPU conversion do not carry unused tensors. Returned mask annotations are unchanged.
Use the existing predicted-IoU decision before expanding candidate masks. Preserve raw unclamped logits for interpolation and the clamped refinement output. Keep the original ordering for refinement and region cleanup.
Copy sparse CUDA batch boundaries together instead of scanning and transferring them once per mask. Limit the boundary table to 256 KiB per mask on average (16384 pairs of int64 indices); retain the original CPU, single-mask and dense paths. This is a conservative crossover, not a universal optimum. Integer RLE counts and ordering remain unchanged.
Convert eager JPEG pixels directly to float32 before division. All uint8 values round identically to the former float64-then-float32 path. Preserve the helper default and async loading arithmetic; normalization is unchanged.
Compute dense positional encoding once within a multi-image prediction call. Pass it to the private prediction helper without keeping persistent cache state, so subsequent model-buffer or autocast changes cannot reuse stale data.
|
Hi @SimBuddy! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
1 participant
Summary
This PR proposes five small changes that reduce unnecessary intermediate work in automatic mask generation, RLE conversion, eager video loading and batched image prediction. Each change is a separate commit with a focused regression test.
Changes
Automatic mask generation
Keep low-resolution logits only while mask-to-mask refinement needs them, then release them before downstream filtering and conversion. Without refinement, avoid storing that unused field altogether.
Use the existing predicted-IoU rejection before full-resolution interpolation. The decoder already supplies the score; candidates that fail the current threshold no longer need expansion. This changes the ordering of an existing decision, with no new threshold or approximation. Interpolation receives raw, unclamped logits. Refinement and region-cleanup configurations retain their existing order.
RLE postprocessing
Batch sparse CUDA boundary transfers instead of scanning and transferring the boundary population once per mask. The batching path requires more than one mask and at most 16,384 boundary changes per mask on average. Each boundary is two int64 indices, so this bounds the copied table to 256 KiB per mask on average.
The original CPU, single-mask and dense paths remain. Unconditional batching regressed on dense populations during validation. The retained limit is a conservative implementation crossover, not a universally optimal threshold.
Video loading
Convert eager JPEG pixels through FP32 directly, avoiding a temporary FP64 representation before copying into FP32 frame storage. All 256 possible byte values match the former FP64-then-FP32 result. The helper default and asynchronous-loading arithmetic are unchanged.
Batched image prediction
Compute dense positional encoding once within a multi-image prediction call and reuse it across the per-image predictions. No persistent cache is introduced; subsequent calls recompute it and can observe model-buffer or autocast changes.
Why
These changes shorten unused tensor lifetimes, apply already-available rejection before expansion, reduce small device transfers, avoid an unnecessary dtype conversion and reuse an image-independent result within its valid lifetime.
Correctness
Eleven focused tests cover mask lifetime/refinement, IoU-filtered result equality and interpolation counts, raw versus clamped logits, exact sparse/dense RLE and boundary guards, byte/JPEG/normalized frame conversion, and batch positional reuse across changed buffers. All pass on CPU/CUDA where applicable, without timing assertions. The exported five-commit series also passes the cumulative tests after each commit on a clean base.
Retained model-level validation includes exact image/AMG annotations, crop/refinement/empty paths, FP32/BF16 image batches, video correction/reverse tracking, distinct 1/2/4-object scenarios and original-versus-patched output/state comparisons within each offload mode. This establishes tested baseline equivalence, not a new segmentation-quality score.
Performance
On the frozen 1080p SAM2.1 tiny AMG benchmark, the combined patch reduced mean runtime from 3499.91 ms to 3100.00 ms (11.43%) and peak PyTorch GPU allocation from 1,903,977,984 to 1,493,497,344 bytes (391.46 MiB less). Results are workload- and hardware-dependent. The comparison used five alternating warm pairs on an RTX 4060 Laptop 8 GB, PyTorch 2.5.1/cu124, CUDA BF16 inference, 32 points/side and 32 points/batch. All 35 returned annotations matched exactly.
A post-cleanup confirmation also matched outputs and allocation, with mean 3671.31→3161.16 ms; its paired timing variation was substantially higher. Both cohorts are retained rather than replacing the original result with the larger percentage.
Individual effects are not additive. No general video FPS or single-image latency gain is claimed.
Testing
Formatter versions: black 24.2.0, usort 1.0.2, ufmt 2.0.0b2, matching the contribution guide. The frozen base has no tracked Python unit-test suite; these tests are new. Modified modules also pass import/compile checks. Heavy benchmark scripts remain outside the production tree.
Scope
Validation used SAM2.1 tiny on one CUDA laptop and a small corpus; percentages are not asserted for other checkpoints, GPUs or configurations. Optional CUDA cleanup could not build in the local toolkit/glibc environment, so both variants used the documented cleanup-skipping installation. The new early-filter path excludes cleanup-active configurations; their real extension behavior was not established locally. Production training and multi-GPU testing were not run. The RLE crossover warrants additional hardware coverage and is isolated in its own commit.
Method
The opportunities were identified and validated using WOB/WOB3 techniques.
WOB is an experimental method of refocussing workloads in data-heavy code. I am stress testing various codebases, to develop the method further, and SAM2 seemed like a worthwhile candidate. The results were good, so I am sharing. Feedback would be greatly appreciated.
WOB is human-designed, human-led, and machine assisted.