Space/time complexity of batched_nms grows quadratically with batch size

[ppwwyyxx]

batched_nms now performs a regular NMS for all boxes combined together. Since NMS takes O(n^2) space and time, the complexity of batch_nms grows quadratically with batch size.

Ideally the complexity should grow linearly with batch size.

[fmassa]

Indeed, we should have a custom implementation which is linear with the batch size.

I believe a naive for-loop implementation in C++ for CPU would be enough, but for CUDA it will probably not be very easy to write.

[NicolasHug]

Here's a benchmark comparing the current batched_nms implementation with a manual for loop, as done in detectron2

Benchmark code

from time import time
from collections import defaultdict

from typing import List
import torch
from torchvision.ops import nms, batched_nms
import pandas as pd


N_BOXES_LIST  = [100, 1_000, 2_000, 3_000, 4_000, 5_000, 6_000, 10_000, 50_000, 100_000]
AVERAGE_OVER = 5


def batched_nms_for_loop(boxes, scores, idxs, iou_threshold):

    result_mask = scores.new_zeros(scores.size(), dtype=torch.bool)
    for id in torch.jit.annotate(List[int], torch.unique(idxs).cpu().tolist()):
        mask = (idxs == id).nonzero(as_tuple=False).view(-1)
        keep = nms(boxes[mask], scores[mask], iou_threshold)
        result_mask[mask[keep]] = True
    keep = result_mask.nonzero(as_tuple=False).view(-1)
    keep = keep[scores[keep].argsort(descending=True)]
    return keep

def bench(func, device):

    iou_threshold = .9

    times = []
    for n_boxes in N_BOXES_LIST:

        boxes = torch.cat((torch.rand(n_boxes, 2), torch.rand(n_boxes, 2) + 10), dim=1).to(device)
        assert max(boxes[:, 0]) < min(boxes[:, 2])  # x1 < x2
        assert max(boxes[:, 1]) < min(boxes[:, 3])  # y1 < y2

        scores = torch.rand(n_boxes).to(device)
        idxs = torch.randint(0, 4, size=(n_boxes,)).to(device)

        tic = time()
        for _ in range(AVERAGE_OVER):
            func(boxes, scores, idxs, iou_threshold)
        times.append((time() - tic) / AVERAGE_OVER)
    return times

d = defaultdict(list)

for func in (batched_nms, batched_nms_for_loop):
    for device in ('cpu', 'cuda'):
        times = bench(func, device)
        d['time'] += times
        d['n_boxes'] += N_BOXES_LIST
        d['device'] += [device] * len(times)
        d['func'] += [func.__name__] * len(times)

print(pd.DataFrame(d).to_csv())

Here are the computation time for different number of boxes, averaged over 5 runs:

	100	1000	2000	3000	4000	5000	6000	10000	50000	100000
('cpu', 'batched_nms')	0.000190401	0.000384903	0.000648355	0.000912523	0.00118756	0.00157652	0.00175991	0.0102302	0.018375	0.0365446
('cpu', 'batched_nms_for_loop')	0.000591469	0.000754929	0.000990105	0.00118504	0.0014338	0.00155549	0.00176287	0.00267587	0.0112071	0.0237366
('cuda', 'batched_nms')	0.000622702	0.000431252	0.000691462	0.0011549	0.00137444	0.0016943	0.00197568	0.00497718	0.258566	1.01611
('cuda', 'batched_nms_for_loop')	0.00184584	0.00196228	0.00204272	0.00250225	0.0026546	0.00274425	0.00283794	0.00383205	0.0320407	0.235898

batched_nms_for_loop is faster than the current version when n_boxes gets bigger. It seems that 6000 boxes is a reasonable threshold to switch from one implementation to the other (both for cpu and gpu). Happy to submit a PR if you agree with the analysis @fmassa.

[fmassa]

@NicolasHug thanks for the benchmark! Can you actually plot the same graph but varying the number of classes? In your example you only have 4 classes, but in many cases we would have much more (up to 80 on COCO for example).
This will give some more representative ideas on how the performance of batched NMS is affected when batching over a larger number of classes

Also, for GPU profiling, don't forget to add torch.cuda.synchronize() before and after measuring the time, otherwise the timings won't be accurate

[NicolasHug]

Thanks for the feedback, here are the time and memory benchmarks with different number of classes. I made sure to call syncrhonize():

main benchmark loop

    for n_boxes in N_BOXES_LIST:

        boxes = torch.cat((torch.rand(n_boxes, 2), torch.rand(n_boxes, 2) + 10), dim=1).to(device)
        assert max(boxes[:, 0]) < min(boxes[:, 2])  # x1 < x2
        assert max(boxes[:, 1]) < min(boxes[:, 3])  # y1 < y2

        scores = torch.rand(n_boxes).to(device)
        idxs = torch.randint(0, n_classes, size=(n_boxes,)).to(device)


        if device == 'cuda':
            torch.cuda.synchronize()
            torch.cuda.empty_cache()
            torch.cuda.reset_peak_memory_stats()
        tic = time()

        for _ in range(AVERAGE_OVER):
            func(boxes, scores, idxs, iou_threshold)

        if device == 'cuda':
            torch.cuda.synchronize()
        times.append((time() - tic) / AVERAGE_OVER)
        mems.append(torch.cuda.max_memory_allocated())

[NicolasHug]

I'm not sure whether the memory info is relevant for CPU, but for GPU it seems that the for-loop version is consistently cheaper memory-wise.

The for-loop version gets faster than the current version when:

• n_boxes > 4k for CPU
• n_boxes > 20k for GPU

Interestingly detectron2 uses torchvision's current version when n_boxes < 40k, but this is for float-precision reasons only.

@fmassa it seems that the for-loop version could at least be useful for CPU when n_boxes > 4k?

[fmassa]

@NicolasHug I agree, we can do a dispatch similar to detectron2 but using a different value depending on CPU / CUDA.

[indigoviolet]

@NicolasHug did you consider making the boxes.numel() threshold (for switching to the for-loop version) a parameter of batched_nms? Seems like that would be useful in the case of having different numbers of classes than this was tuned for, or for GPU (per your comment about using a higher threshold for GPU).

Alternatively, if you expose batched_nms_coordinate_trick, users could call that directly when appropriate.

[NicolasHug]

I think the reasoning behind using the same threshold for GPU is that we don't expect batched_nms to be used with much more than 4k boxes anyway (the current threshold). Do you have a use-case where you have more than 4k boxes?

did you consider making the boxes.numel() threshold

That might be reasonable, although I'm not sure we usually expose that kind of implementation detail to users. @fmassa WDYT?

[fmassa]

I would prefer to leave this threshold as an implementation detail. In the future we can come up with different implementations that might completely make this parameter not relevant anymore.

Additionally, I would prefer not to expose batched_nms_coordinate_trick, as again it's an implementation detail.

Do you have an use-case in which the new implementation (with a single threshold) is significantly slower than it was before?