Improved performance in torchvision backend

[ZephyrKeXiner]

Hi maintainer @onuralpszr,

While benchmarking the postprocess backends, I noticed a potential bottleneck in the torchvision backend.

In this step, the backend computes a dense N x N IoU/IoS matrix on GPU and then copies the whole matrix back to CPU. For large numbers of boxes, this transfer can become expensive. For example, 10,000 x 10,000 float32 values require about 400 MB to transfer.

Instead of transferring the full dense matrix, I tried using a packed bitmask representation. The Triton implementation computes the thresholded IOS matches directly and returns a packed uint32 bitmask, which is then consumed by the greedy merge loop on CPU.

On a Tesla T4, this gives a clear speedup for large GreedyNMM + IOS cases:

boxes	numpy	torchvision matrix	torch bool mask	Triton packed bitmask	speedup vs matrix	speedup vs torch mask
1,000	41.6 ms	10.4 ms	4.1 ms	8.7 ms	1.2x	0.5x
2,000	91.6 ms	30.7 ms	14.8 ms	10.0 ms	3.1x	1.5x
5,000	587.8 ms	173.1 ms	40.4 ms	29.3 ms	5.9x	1.4x
10,000	2249.4 ms	618.5 ms	188.9 ms	73.1 ms	8.5x	2.6x

Parity with the existing NumPy result is preserved in the benchmark.

The Triton packed-bitmask path is not beneficial for small inputs, but becomes the fastest option from around 2k boxes and gives up to 8.5x speedup over the current torchvision matrix path at 10k boxes.The median speedup is significant, although the Triton path currently shows higher variance for larger inputs. Even at p90, it remains faster than the dense matrix path.

[ZephyrKeXiner]

Torch boolean mask path

The current torchvision backend computes a dense N x N float32 IoU/IoS matrix on GPU and then copies the whole matrix back to CPU before running the greedy merge loop.

For GreedyNMM, the CPU merge loop does not actually need the exact IoU/IoS value after the threshold check. It only needs to know whether each pair is a match:

matrix[i, j] >= match_threshold

Triton packed bitmask path

The Triton path goes one step further than the torch boolean-mask path.

Instead of materializing a dense float32 matrix or a dense boolean matrix, the Triton kernel computes the thresholded IOS match result directly and packs the result into integer bitsets.

Each uint32 word stores 32 pairwise match results:

1 bit = whether candidate j matches keeper i

[onuralpszr]

@ZephyrKeXiner hey, can you open PR so I can check and test as well, it looks promising in first glance, but let me also test as well plus if you have dense example that would be also cool for test purpose as well (plus model if needed)

[ZephyrKeXiner]

Before open a PR, I want to review the implementation and benchmark one more time to make sure the conclusion is precise. The current result looks promising, but I want to be careful about the scope: this optimization mainly targets large GreedyNMM + IOS postprocess inputs, not all postprocess cases.

I will open the draft PR after the final check.

[onuralpszr]

Even it if it is wrong a PR would be fine because we can discuss and test and see if not we can close PR if all good we can merge it is a research process as well so all good

[ZephyrKeXiner]

Thank you very much. I will open a draft PR so we can discuss and test it there.

[steezbert]

Hi all,

I have a quick question/discovery that might be related to this discussion/issue.

I am mainly using the sahi.postprocess.combine module and noticed some differences after switching from SAHI v0.11.36 to v0.12.1.

I am combining 33337 labels from 9866 images using NMM.
Using v0.11 the label reduction took 5sec on my machine.
Using v0.12 with the torchvision backend it runs into a memory error:

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 4.14 GiB. GPU 0 has a total capacity of 3.68 GiB of which 3.58 GiB is free.

Using v0.12 with the numpy backend the label reduction took 35sec.

So I am wondering if this issue is related to the topic, why there is a performance downgrade from v0.11 to v0.12 and if you are able/planning to fix the issue.

Just wanted to let you know about my findings. If you have further questions let me know.

[onuralpszr]

@steezbert do you have minimal repro code so I can check or data so we can try as well. ALso you can switch to use numpy as a backend if you like but you said "numpy took 35 secs" what was before 0.12 ? that would be nice info as well.

[steezbert]

So, this is the reduced code snipped i am using:

import time

import geopandas
import numpy
from sahi.postprocess.combine import NMMPostprocess
from sahi.prediction import ObjectPrediction
from shapely import box

# SAHI v0.12 only
# from sahi.postprocess.backends import set_postprocess_backend
# _ = set_postprocess_backend("numpy")


def reduce_overlapping_labels(
    label_polygons: geopandas.GeoDataFrame,
    label_column: str = "yolo_label",
    score_column: str = "yolo_confidence",
    threshold: float = 0.5,
    metric: str = "IOU",
    class_agnostic: bool = True,
):
    # convert label_polygons into list of ObjectPrediction objects
    obj_pred_list = [
        ObjectPrediction(
            bbox=numpy.array(row.geometry.bounds),
            category_id=int(row[label_column]),
            score=row[score_column],
        )
        for _idx, row in label_polygons.iterrows()
    ]

    # set postprocess based on selected method
    postprocess = NMMPostprocess(
        match_threshold=threshold,
        match_metric=metric,
        class_agnostic=class_agnostic,
    )

    # postprocess obj_pred_list
    obj_pred_list = postprocess(obj_pred_list)

    # convert back to separat tuples for polygons, labels and confs
    polygons, labels, confs = zip(
        *[(box(*(obj.bbox.to_xyxy())), str(obj.category.id), float(obj.score.value)) for obj in obj_pred_list],
        strict=True,
    )

    # convert to GeoDataFrame
    reduced_label_polygons = geopandas.GeoDataFrame(
        data={"yolo_label": list(labels), "yolo_confidence": list(confs)},
        geometry=list(polygons),
        crs=label_polygons.crs,
    )

    return reduced_label_polygons


label_polygons = geopandas.read_file("~/YOLO_Label_Example.geojson")

t1 = time.time()

results = reduce_overlapping_labels(
    label_polygons=label_polygons,
    label_column="yolo_label",
    score_column="yolo_confidence",
    threshold=0.3,
    metric="IOS",
    class_agnostic=True,
)

t2 = time.time()
print(f"reducing overlapping labels: {t2 - t1}")

This is the example file YOLO_Label_Example.zip
(zipped geojson, let me know if you want a different format, this is the easiest with geopandas, unfortunately github doesn't allow .geojson uploads)

sahi version	backend	runtime
v0.11	auto (which is torch?)	~5sec
v0.12	numpy	~35sec
v0.12	torchvision	memory error

Thank you for looking into the issue. Let me know if you need further information.

[onuralpszr]

The quadratic allocation is fixed in 0.12.2, just released: https://github.com/obss/sahi/releases/tag/0.12.2

@steezbert your report was the same root cause, so your case is covered too.

What changed

@ZephyrKeXiner your diagnosis was right, and it went further than the transfer cost. The greedy loops only ever use matrix >= match_threshold, never the metric values, so the dense N x N matrix never needed to exist on either side of the bus. 0.11 knew this implicitly: it used a shapely STRtree and only compared boxes the index reported as nearby. 0.12 replaced that with the dense matrix, which is O(N^2) in time and memory whatever the layout, and nmm_from_matrix added five more N x N bool arrays on top.

0.12.2 builds the thresholded adjacency directly from the pairs an STRtree reports as intersecting and stores it as CSR. NumPy, TorchVision and Numba all route large inputs there. Below 2000 boxes everything stays dense, where it is faster, as does any non-positive threshold, since that matches disjoint pairs an intersection prefilter never enumerates.

Numbers

@steezbert's 33337 boxes, IOS, threshold 0.3, on an i7-13850HX with an RTX 4000 Ada Laptop (12 GB). Outputs are unchanged: 6571 predictions for greedy_nmm, 4639 for nmm.

CPU:

backend	greedy_nmm before	after	nmm before	after
numpy	15.88s	0.18s	27.06s	0.21s
torchvision	6.22s	0.18s	17.34s	0.21s
numba	1.53s	1.55s	18.56s	0.22s

CUDA:

backend	greedy_nmm before	after	nmm before	after
numpy	15.93s	0.26s	27.00s	0.21s
torchvision	OOM, 4.14 GiB	0.17s	OOM, 4.14 GiB	0.23s
numba	1.09s	1.09s	18.82s	0.22s

The OOM reproduced exactly as reported before the fix, at the same 4.14 GiB.

Where this leaves the Triton idea

Worth being straight about: your benchmark measured Triton against the dense matrix path, and that path no longer runs for the sizes where Triton won. At 10000 boxes you had 618.5ms for the dense matrix and 73.1ms for the packed bitmask. The sparse path is now in that same range on CPU alone, so the comparison needs redoing against 0.12.2 before we can say what Triton buys.

There is a real gap it could still fill. Above the cutoff the TorchVision backend now delegates to the NumPy sparse functions, so it does not touch the GPU at all for large inputs. A kernel that produces the thresholded adjacency on-device, in the packed form you described, would put the GPU back to work without ever materializing a dense matrix, and the CSR consumer is already in place for it to feed.

Two things to watch if you pick it up. The STRtree build is now the dominant cost for nms and greedy_nmm (around 156ms of the 33337-box run), so that is the number to beat, not the old dense timing. And a GPU kernel computing in float32 against a CPU path at the input dtype can disagree on pairs sitting exactly on the threshold, so parity checks should include boundary cases rather than random boxes only.

A PR is still very welcome, benchmarked against 0.12.2. you might update your PR @ZephyrKeXiner so we can take look again.

[ZephyrKeXiner]

This is an excellent fix! I’m genuinely impressed by how you stepped back from optimizing the dense path and eliminated the need for the dense matrix altogether.

I was too focused on making the existing O(N^2) path faster and didn’t question deeply enough whether that matrix needed to exist in the first place. Sorry that my thinking didn’t go far enough there, and thank you for the detailed explanation and benchmarks.

[ZephyrKeXiner]

Hi @onuralpszr,

First of all, thank you for the elegant solution to the OOM issue. The sparse backend avoids materializing a dense O(N²) overlap matrix for large inputs.

I noticed that when there are at least 2,000 bounding boxes, post-processing paths other than NMS with IoU fall back to the CPU-based sparse implementation when match_threshold > 0.

I have not benchmarked the runtime cost yet, but I am wondering whether a GPU-friendly sparse or tiled algorithm could retain the memory benefits while improving performance on GPU-equipped systems. Do you think such an optional GPU backend would fit SAHI’s design?

[onuralpszr]

I did some changes can you check latest release notes please.

[ZephyrKeXiner]

Thank you for pointing me to the 0.12.5 release notes. I have just read through them and the related PRs, and I am genuinely impressed by both the speed of your response and the accuracy with which you identified the underlying problem.

Recognizing that the sparse pair list could still approach O(N²) in crowded layouts, and replacing it with on-demand STRtree queries, is a very elegant solution. The density-aware routing and benchmark results are also impressive.

I really respect the way you analyze the actual bottleneck, validate it quantitatively, and turn it into a practical implementation so quickly. There is a lot for me to learn from your approach. Thank you again for the excellent work and for sharing the details so openly.❤️

[onuralpszr]

@ZephyrKeXiner thank you for kind answer and analyse of yours. Also we can still experiment maybe "just thinking" DALI + Trition + TensorRT paths to make it faster but experiment required.

[onuralpszr]

Another round is merged I did some docs clean ups as well. Still experimenting for more cases and benchmark we need to make as well. But thank you for your work.