Performance: Dense `O(n×m)` IoU matrices cause excessive memory usage for large-scale evaluations

[musaqlain]

_overlap_all() allocates two massive, dense (n_truth × n_pred) NumPy arrays, even though the STRtree query immediately before it proves that only a tiny fraction of pairs actually intersect.

https://github.com/weecology/DeepForest/blob/main/src/deepforest/IoU.py#L54-L58

These dense matrices are then passed to linear_sum_assignment() and used to compute a third dense iou_mat, further tripling the memory cost.

https://github.com/weecology/DeepForest/blob/main/src/deepforest/IoU.py#L115

https://github.com/weecology/DeepForest/blob/main/src/deepforest/IoU.py#L122-L128

Impact

For a realistic large-tile evaluation with 5,000 ground truth boxes and 10,000 predictions, this allocates:

• intersections: 5,000 × 10,000 × 8 bytes = ~381 MB
• unions: same = ~381 MB
• iou_mat (line 123): same = ~381 MB
• np.zeros_like (line 126): same = ~381 MB
• Total: ~1.5 GB for a single image evaluation

In reality, only a tiny fraction of these pairs actually intersect (typically <1%), so >99% of that memory stores zeros. _overlap_all is called per image via evaluate_boxes → evaluate_image_boxes → compute_IoU, so this becomes a bottleneck when evaluating large NEON plots or aerial survey tiles containing thousands of trees.

Steps to Reproduce

Here is a standalone script that measures the memory waste using tracemalloc:

import time
import tracemalloc
import numpy as np
import geopandas as gpd
from shapely.geometry import box
from deepforest.IoU import _overlap_all

def make_dummy_data(n, img_size=10000, box_size=50):
    xs = np.random.uniform(0, img_size - box_size, n)
    ys = np.random.uniform(0, img_size - box_size, n)
    geoms = [box(x, y, x + box_size, y + box_size) for x, y in zip(xs, ys)]
    return gpd.GeoDataFrame({"score": np.random.uniform(0.5, 1.0, n)}, geometry=geoms)

np.random.seed(42)
truth = make_dummy_data(5000)
preds = make_dummy_data(10000)

tracemalloc.start()
t0 = time.time()

inter, _, _, _ = _overlap_all(preds, truth)

t1 = time.time()
_, peak = tracemalloc.get_traced_memory()
tracemalloc.stop()

sparsity = np.count_nonzero(inter) / inter.size
print(f"Shape: {inter.shape}")
print(f"Sparsity: {sparsity:.4f}")
print(f"Peak memory: {peak / 1024 / 1024:.1f} MB")
print(f"Time: {t1-t0:.2f}s")

#Proposed Fix
-Use scipy.sparse.csr_matrix instead of np.zeros() to store the intersections and unions, reducing the spatial complexity from O(N*M) to O(k) (where k is the number of actual overlapping pairs).
-For linear_sum_assignment which requires a dense matrix, extract only the dense sub-matrix consisting of rows and columns that have at least one intersection, then map the assignment indices back to the original arrays. This greatly reduces the problem size for the cubic-time Hungarian algorithm.

[musaqlain]

Since the STRtree already identifies the sparse set of overlapping pairs, the function should use a sparse representation instead of dense matrices. This would reduce memory from O(n×m) to O(M) where M is the number of actual intersections (typically << n×m).

[vickysharma-prog]

Interesting analysis! One thing to consider: scipy.optimize.linear_sum_assignment requires a dense matrix, so the sparse approach would need to extract a reduced sub-matrix of rows/columns with actual intersections, then map the assignment indices back to the original arrays. The index bookkeeping adds complexity but the memory savings could be significant for large-scale evaluations.

[ethanwhite]

@henrykironde has been working on memory issues so I'm going to have him take a look at this issue. If this is something different from what he's already addressing and doesn't cause issues with that other work he can assign @musaqlain to work on it.

[jveitchmichaelis]

Thanks @musaqlain. If you're sharing sample/trace code, please could you also execute and share the results so we can have a look?

We tend to not do anything in parallel here, unless running on multiple GPUs. Although this is a large allocation, it should be a fixed overhead for a dataset with fixed size images. I agree this is a sensible thing to improve, but I wouldn't expect this to be the cause of CPU OOMs - a couple GB on the cluster is not our concern I suspect (@henrykironde?). I would argue we care more about speed of execution here, as this code gets called a lot during training and on larger validation datasets it's significant.

For a realistic large-tile evaluation with 5,000 ground truth boxes and 10,000 predictions, this allocates:

For large images, we would probably optimize slightly differently and use tiling + reduction over the tiles? For example this would assume that either we're predicting an enormous image in one go, or we're storing the results and then comparing against ground truth. It's certainly a use-case we should think about, but more often we only run inference on images that big, not evaluation.

[musaqlain]

Thanks @jveitchmichaelis, here are the executed results:

Creating 5000 ground truth and 10000 prediction boxes...

Shape: (5000, 10000)
Sparsity: 0.0001
Peak memory: 763.9 MB
Time: 0.97s

I came across this while studying the evaluation pipeline for point detection (#809), specifically PR #1182, where keypoint matching uses Euclidean distance instead of IoU but follows the same cost-matrix → linear_sum_assignment pattern. That led me to look at how the current IoU matching works.

What caught my attention is that _overlap_all() already computes the sparse overlap pairs via STRtree (line 55: p_idx, t_idx), but then discards that sparse knowledge by filling it into dense (n_truth × n_pred) matrices. The dense matrices are then passed to linear_sum_assignment which is O(n²m) — reducing the input dimensions to only the rows/columns with at least one intersection would give a direct speed win, which I agree is the bigger concern here.

Happy to wait on Henry assessment of whether it works. thans 🙌

[jveitchmichaelis]

I agree, this seems like a sensible thing to do. I'd suggest opening a PR with the modification and we can take a look from there?

I would actually go simpler. Just run a filter on the output from SRTree and identify any matches that are 1:1 already, pop those separately with their IoUs, and then run Hungarian on the remainder (if any). There's probably no need to go for a sparse matrix here, that second pool of points should already be very small.

Bear in mind that RetinaNet already performs non-max suppression inside the model class (we don't control that), so the evaluation/matching stage is already working with that "cleaned" output. Slightly different when this is used inside the training loop (e.g. DETR) and all predictions go into Hungarian to force the model to not predict overlapping boxes to begin with.

[jveitchmichaelis]

The memory penalty might still be worth it to take advantage of vectorization, though you could definitely look at returning sparse arrays for intersection + union.

EDIT: Looking briefly, we would filter out p_idx where len(t_idx) == 1?

[bw4sz]

I agree with @jveitchmichaelis "I would actually go simpler. Just run a filter on the output from SRTree and identify any matches that are 1:1 already, pop those separately with their IoUs, and then run Hungarian on the remainder (if any). There's probably no need to go for a sparse matrix here, that second pool of points should already be very small." I would support a PR in this area.

[vickysharma-prog]

Quick thought on the 1:1 filter - two preds could both match the same truth box with len(t_idx)==1, so might need to check both directions. Happy to help with a PR if @musaqlain is busy!

[musaqlain]

I agree with @jveitchmichaelis "......." I would support a PR in this area.

sure, I will raise PR accordingly