⚡️ Speed up method SAM2Model._apply_non_overlapping_constraints by 95%

[codeflash-ai]

📄 95% (0.95x) speedup for SAM2Model._apply_non_overlapping_constraints in ultralytics/models/sam/modules/sam.py

⏱️ Runtime : 3.12 milliseconds → 1.60 milliseconds (best of 66 runs)

📝 Explanation and details

The optimized code achieves a 95% speedup by making two key optimizations to the _apply_non_overlapping_constraints method:

Key Optimizations:

1. Replaced torch.argmax with torch.max: Changed from torch.argmax(pred_masks, dim=0, keepdim=True) to _, max_obj_inds = torch.max(pred_masks, dim=0, keepdim=True). This is faster because torch.max returns both values and indices in a single fused operation, while torch.argmax only computes indices but still needs to traverse the data similarly.

2. Pre-computed clamped tensor: Moved torch.clamp(pred_masks, max=-10.0) outside the torch.where call by pre-computing it as min_mask. This avoids redundant tensor clamping operations within the conditional assignment.

Performance Impact:

• Line profiler shows the critical torch.argmax line dropped from 2.66ms to 1.19ms (55% reduction)
• The torch.where operation became more efficient, dropping from 764μs to 422μs
• Overall function runtime improved from 3.12ms to 1.60ms

Test Case Performance:
The optimization shows particularly strong gains on larger tensors:

• Large spatial shapes (128x128): 404% faster
• Large batch sizes (16 objects): 99% faster
• Smaller tensors show modest but consistent improvements

Why This Works:
In PyTorch, torch.max is optimized as a single kernel operation that finds both maximum values and their indices simultaneously, while torch.argmax performs similar work but discards the values. Pre-computing the clamped tensor reduces redundant memory operations in the conditional assignment. These optimizations are especially effective for the multi-object segmentation use case where batch sizes and spatial dimensions are typically large.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 31 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import pytest
import torch
from ultralytics.models.sam.modules.sam import SAM2Model

# ------------------ Unit Tests ------------------

# 1. Basic Test Cases


def test_single_object_mask_is_unchanged():
    # Single object, should return the same mask
    mask = torch.tensor([[[[0.1, 0.2], [0.3, 0.4]]]])
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 1.60μs -> 1.59μs (0.377% faster)


def test_two_objects_non_overlapping():
    # Two objects, no overlap (each object has max at different locations)
    mask = torch.tensor(
        [
            [[[0.9, 0.1], [0.8, 0.2]]],  # object 0
            [[[0.1, 0.8], [0.2, 0.9]]],  # object 1
        ]
    )
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 56.1μs -> 57.7μs (2.86% slower)
    # object 0 should keep (0,0) and (1,0); object 1 keeps (0,1) and (1,1)
    expected = torch.tensor([[[[0.9, -10.0], [0.8, -10.0]]], [[[-10.0, 0.8], [-10.0, 0.9]]]])


def test_two_objects_full_overlap():
    # Two objects, both have max at all locations, but object 1 always higher
    mask = torch.tensor([[[[0.5, 0.5], [0.5, 0.5]]], [[[0.7, 0.7], [0.7, 0.7]]]])
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 42.0μs -> 42.1μs (0.183% slower)
    # Only object 1 keeps its values, object 0 gets clamped
    expected = torch.tensor([[[[-10.0, -10.0], [-10.0, -10.0]]], [[[0.7, 0.7], [0.7, 0.7]]]])


def test_three_objects_partial_overlap():
    # Three objects, partial overlap
    mask = torch.tensor(
        [
            [[[0.9, 0.2], [0.1, 0.4]]],  # obj 0: max at (0,0)
            [[[0.2, 0.9], [0.3, 0.1]]],  # obj 1: max at (0,1)
            [[[0.1, 0.3], [0.9, 0.8]]],  # obj 2: max at (1,0) and (1,1)
        ]
    )
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 40.8μs -> 41.7μs (2.10% slower)
    expected = torch.tensor(
        [[[[0.9, -10.0], [-10.0, -10.0]]], [[[-10.0, 0.9], [-10.0, -10.0]]], [[[-10.0, -10.0], [0.9, 0.8]]]]
    )


# 2. Edge Test Cases


def test_all_objects_equal_scores():
    # All objects have the same score at each location, only one object (lowest index) keeps value
    mask = torch.ones((3, 2, 2, 2))
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 41.2μs -> 40.8μs (0.993% faster)
    # Only the first object keeps the value due to argmax returning first max index
    expected = torch.cat([torch.ones(1, 2, 2, 2), -10.0 * torch.ones(2, 2, 2, 2)], dim=0)


def test_negative_and_positive_scores():
    # Some values negative, some positive
    mask = torch.tensor([[[[0.5, -1.0], [0.2, -0.2]]], [[[-0.5, 0.7], [-0.3, 0.9]]]])
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 42.6μs -> 41.4μs (3.01% faster)
    # (0,0): obj0 keeps, (0,1): obj1 keeps, (1,0): obj0 keeps, (1,1): obj1 keeps
    expected = torch.tensor([[[[0.5, -10.0], [0.2, -10.0]]], [[[-10.0, 0.7], [-10.0, 0.9]]]])


def test_large_negative_scores():
    # All values are already less than -10.0
    mask = torch.full((2, 2, 2, 2), -20.0)
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 39.0μs -> 38.3μs (1.89% faster)


def test_single_pixel_masks():
    # Masks of shape (N, 1, 1, 1)
    mask = torch.tensor([[[[-0.1]]], [[[0.2]]], [[[0.0]]]])
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 39.4μs -> 38.6μs (2.16% faster)
    # Only object 1 keeps value
    expected = torch.tensor([[[[-10.0]]], [[[0.2]]], [[[-10.0]]]])


def test_device_consistency_cpu():
    # Should work on CPU
    mask = torch.randn(2, 2, 2, 2)
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 40.1μs -> 40.5μs (0.992% slower)


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_input_not_modified_inplace():
    # The input tensor should not be modified in-place
    mask = torch.tensor([[[[0.5, 0.1], [0.2, 0.3]]], [[[0.1, 0.9], [0.2, 0.4]]]])
    mask_clone = mask.clone()
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    _ = codeflash_output  # 48.4μs -> 49.6μs (2.44% slower)


def test_batch_size_one_and_multi():
    # Batch size 1 should return unchanged, batch size >1 should change
    mask1 = torch.rand(1, 2, 2, 2)
    mask2 = torch.rand(2, 2, 2, 2)
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask1)
    out1 = codeflash_output  # 1.33μs -> 1.35μs (1.48% slower)
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask2)
    out2 = codeflash_output  # 38.8μs -> 38.0μs (1.99% faster)


# 3. Large Scale Test Cases


def test_large_batch_and_spatial_shape():
    # Large batch and spatial dims, but under 100MB
    B, C, H, W = 16, 1, 32, 32  # 16*1*32*32*4 = 65,536 bytes
    mask = torch.rand(B, C, H, W)
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 216μs -> 109μs (99.1% faster)
    # Check that for each location, only one object keeps the value
    max_inds = torch.argmax(mask, dim=0, keepdim=True)
    for b in range(B):
        # For each batch, only keep if b == max_inds at that location
        keep = max_inds == b
        # All other locations should be clamped
        suppressed = out[b][~keep.squeeze(0)]


def test_large_spatial_shape():
    # Large spatial dimensions, but under 100MB
    B, C, H, W = 2, 1, 128, 128  # 2*1*128*128*4 = 131,072 bytes
    mask = torch.rand(B, C, H, W)
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 1.55ms -> 307μs (404% faster)
    # For each pixel, only one batch should have non-suppressed value
    max_inds = torch.argmax(mask, dim=0, keepdim=True)
    for i in range(H):
        for j in range(W):
            # Only one batch should be non-clamped
            kept = [out[b, 0, i, j].item() > -10.0 for b in range(B)]


def test_large_batch_size():
    # Large batch size, small spatial
    B, C, H, W = 32, 1, 4, 4  # 32*1*4*4*4 = 2,048 bytes
    mask = torch.rand(B, C, H, W)
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 57.4μs -> 56.7μs (1.21% faster)
    # For each pixel, only one batch should have non-suppressed value
    max_inds = torch.argmax(mask, dim=0, keepdim=True)
    for i in range(H):
        for j in range(W):
            kept = [out[b, 0, i, j].item() > -10.0 for b in range(B)]


def test_large_channel_dim():
    # Large channel dimension, should not affect logic (channels are not involved in argmax)
    B, C, H, W = 3, 8, 4, 4
    mask = torch.rand(B, C, H, W)
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 61.6μs -> 54.2μs (13.5% faster)
    # For each channel and pixel, only one batch should have non-suppressed value
    max_inds = torch.argmax(mask, dim=0, keepdim=True)
    for c in range(C):
        for i in range(H):
            for j in range(W):
                kept = [out[b, c, i, j].item() > -10.0 for b in range(B)]


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import torch
from ultralytics.models.sam.modules.sam import SAM2Model

# unit tests

# ------------------- BASIC TEST CASES -------------------


def test_single_object_no_constraint():
    # Single object: output should be unchanged
    mask = torch.tensor([[[[0.1, 0.2], [0.3, 0.4]]]])
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    result = codeflash_output  # 1.42μs -> 1.41μs (0.141% faster)


def test_two_objects_non_overlapping():
    # Two objects, each dominates a different pixel
    mask = torch.tensor([[[[0.9, 0.1], [0.1, 0.8]]], [[[0.1, 0.8], [0.9, 0.2]]]])
    # pixel [0,0]: obj0, [0,1]: obj1, [1,0]: obj1, [1,1]: obj0
    expected = torch.tensor([[[[0.9, -10.0], [-10.0, 0.8]]], [[[-10.0, 0.8], [0.9, -10.0]]]])
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    result = codeflash_output  # 54.0μs -> 56.5μs (4.40% slower)


def test_two_objects_fully_overlapping():
    # Both objects have same values everywhere, so only object 0 should remain (argmax returns first occurrence)
    mask = torch.ones(2, 1, 2, 2)
    expected = torch.tensor([[[[1.0, 1.0], [1.0, 1.0]]], [[[-10.0, -10.0], [-10.0, -10.0]]]])
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    result = codeflash_output  # 40.9μs -> 41.3μs (0.942% slower)


def test_three_objects_partial_overlap():
    # Three objects, each pixel is dominated by a different object
    mask = torch.tensor([[[[0.9, 0.1], [0.1, 0.2]]], [[[0.2, 0.8], [0.1, 0.2]]], [[[0.1, 0.2], [0.7, 0.9]]]])
    # [0,0]: 0, [0,1]: 1, [1,0]: 2, [1,1]: 2
    expected = torch.tensor(
        [[[[0.9, -10.0], [-10.0, -10.0]]], [[[-10.0, 0.8], [-10.0, -10.0]]], [[[-10.0, -10.0], [0.7, 0.9]]]]
    )
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    result = codeflash_output  # 40.5μs -> 40.3μs (0.568% faster)


def test_negative_and_large_values():
    # Negative and large values: suppression should clamp to -10.0, not below
    mask = torch.tensor([[[[20.0, -5.0], [0.0, 0.0]]], [[[10.0, 30.0], [-20.0, 0.0]]]])
    # [0,0]: 0, [0,1]: 1, [1,0]: 0, [1,1]: 0 (tie, first wins)
    expected = torch.tensor([[[[20.0, -10.0], [0.0, 0.0]]], [[[-10.0, 30.0], [-10.0, -10.0]]]])
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    result = codeflash_output  # 39.4μs -> 40.3μs (2.21% slower)


# ------------------- EDGE TEST CASES -------------------


def test_all_zeros():
    # All zeros: only first object should remain
    mask = torch.zeros(3, 1, 2, 2)
    expected = torch.zeros(3, 1, 2, 2)
    expected[1:] = -10.0
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    result = codeflash_output  # 35.8μs -> 35.6μs (0.499% faster)


def test_all_negative():
    # All negative values, but different
    mask = torch.tensor([[[[-1.0, -2.0], [-3.0, -4.0]]], [[[-5.0, -1.5], [-2.5, -3.5]]]])
    # [0,0]: 0, [0,1]: 1, [1,0]: 0, [1,1]: 0
    expected = torch.tensor([[[[-1.0, -10.0], [-3.0, -4.0]]], [[[-10.0, -1.5], [-10.0, -10.0]]]])
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    result = codeflash_output  # 41.4μs -> 39.7μs (4.25% faster)


def test_minimal_shape():
    # Minimal shape: (1, 1, 1, 1)
    mask = torch.tensor([[[[0.5]]]])
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    result = codeflash_output  # 1.33μs -> 1.37μs (3.28% slower)


def test_large_negative_values():
    # Values already below -10.0 should remain unchanged for suppressed
    mask = torch.tensor([[[[0.5, -20.0], [0.5, 0.5]]], [[[0.1, 0.6], [-20.0, 0.2]]]])
    # [0,0]: 0, [0,1]: 1, [1,0]: 0, [1,1]: 0
    expected = torch.tensor([[[[0.5, -10.0], [0.5, 0.5]]], [[[-10.0, 0.6], [-10.0, -10.0]]]])
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    result = codeflash_output  # 61.1μs -> 63.2μs (3.29% slower)


def test_different_channel_dimension():
    # Test with channel dimension > 1
    mask = torch.tensor(
        [[[[0.1, 0.2], [0.3, 0.4]], [[0.5, 0.6], [0.7, 0.8]]], [[[0.9, 0.8], [0.7, 0.6]], [[0.5, 0.4], [0.3, 0.2]]]]
    )  # shape (2, 2, 2, 2)
    # For each pixel and channel, check which batch has max
    expected = torch.tensor(
        [
            [[[0.1, 0.2], [0.3, 0.4]], [[0.5, 0.6], [0.7, 0.8]]],
            [[[0.9, 0.8], [0.7, 0.6]], [[-10.0, -10.0], [-10.0, -10.0]]],
        ]
    )
    # For channel 0, all pixels: batch 1 wins; for channel 1, batch 0 wins
    expected[0, 0, :, :] = torch.tensor([-10.0, -10.0, -10.0, -10.0]).reshape(2, 2)
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    result = codeflash_output  # 39.0μs -> 39.6μs (1.67% slower)


# ------------------- LARGE SCALE TEST CASES -------------------


def test_large_batch_and_spatial():
    # Large batch and spatial size, but <100MB
    batch = 8
    H, W = 32, 32  # 8*1*32*32*4 = 32KB
    mask = torch.rand(batch, 1, H, W)
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    result = codeflash_output  # 177μs -> 86.3μs (106% faster)
    # For each pixel, only one batch should have a non-suppressed value
    for i in range(H):
        for j in range(W):
            vals = result[:, 0, i, j]


def test_large_channel():
    # Large channel dimension
    mask = torch.rand(4, 16, 8, 8)
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    result = codeflash_output  # 154μs -> 75.3μs (106% faster)
    # For each pixel and channel, only one batch should have non-suppressed value
    for c in range(16):
        for i in range(8):
            for j in range(8):
                vals = result[:, c, i, j]


def test_gradient_preservation():
    # Ensure gradients flow through the kept values
    mask = torch.tensor([[[[0.5, 0.1], [0.2, 0.3]]], [[[0.1, 0.6], [0.2, 0.3]]]], requires_grad=True)
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    out = codeflash_output  # 56.4μs -> 61.0μs (7.46% slower)
    loss = out.sum()
    loss.backward()
    # Only the kept values should have nonzero gradient
    grads = mask.grad


def test_device_support():
    # Should work on CUDA if available
    if torch.cuda.is_available():
        mask = torch.rand(3, 1, 4, 4, device="cuda")
        codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
        result = codeflash_output


def test_non_contiguous_input():
    # Should work with non-contiguous tensors
    mask = torch.rand(4, 1, 8, 8).transpose(0, 2)
    codeflash_output = SAM2Model._apply_non_overlapping_constraints(mask)
    result = codeflash_output  # 58.7μs -> 58.1μs (1.01% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-SAM2Model._apply_non_overlapping_constraints-mirdk9al and push.