Repository: pytorch/vision
Pull request: [OPS, IMP] New batched_nms implementation
Closes https://github.com/pytorch/vision/issues/1311

This PR introduces a new implementation of `batched_nms`, which is faster than the current one in some cases (refer to benchmarks in issue)

Problem statement:
Space/time complexity of batched_nms grows quadratically with batch size
`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.

Hints from issue discussion:
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.
Here's a benchmark comparing the current `batched_nms` implementation with a manual for loop, as done in [detectron2](https://github.com/facebookresearch/detectron2/blob/d6e8c1309b93068c97f89f28bba09190ccc4da6a/detectron2/layers/nms.py#L32)


<details>
<summary>Benchmark code</summary>

```py
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())
```

</details>

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  |

![image](https://user-images.githubusercontent.com/1190450/108357988-11347780-71e6-11eb-9b8e-2f95e9728111.png)


`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.
@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
Thanks for the feedback, here are the time and memory benchmarks with different number of classes. I made sure to call `syncrhonize()`:


![image](https://user-images.githubusercontent.com/1190450/108499213-dd208b80-72a5-11eb-8179-2c2b8b179fe5.png)



<details>
<summary>main benchmark loop </summary>

```py
    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())
```

</details>
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](https://github.com/facebookresearch/detectron2/blob/d6e8c1309b93068c97f89f28bba09190ccc4da6a/detectron2/layers/nms.py#L28) 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?
@NicolasHug I agree, we can do a dispatch similar to `detectron2` but using a different value depending on CPU / CUDA.

Commit message: Use same threshold (4k) for CPU and GPU

Patch to review:
diff --git a/test/test_ops.py b/test/test_ops.py
index 244960cebcc..fd4b356bb69 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -461,6 +461,24 @@ def test_nms_cuda_float16(self):
         keep16 = ops.nms(boxes.to(torch.float16), scores.to(torch.float16), iou_thres)
         self.assertTrue(torch.all(torch.eq(keep32, keep16)))
 
+    def test_batched_nms_implementations(self):
+        """Make sure that both implementations of batched_nms yield identical results"""
+
+        num_boxes = 1000
+        iou_threshold = .9
+
+        boxes = torch.cat((torch.rand(num_boxes, 2), torch.rand(num_boxes, 2) + 10), dim=1)
+        assert max(boxes[:, 0]) < min(boxes[:, 2])  # x1 < x2
+        assert max(boxes[:, 1]) < min(boxes[:, 3])  # y1 < y2
+
+        scores = torch.rand(num_boxes)
+        idxs = torch.randint(0, 4, size=(num_boxes,))
+        keep_vanilla = ops.boxes._batched_nms_vanilla(boxes, scores, idxs, iou_threshold)
+        keep_trick = ops.boxes._batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold)
+
+        err_msg = "The vanilla and the trick implementation yield different nms outputs."
+        self.assertTrue(torch.allclose(keep_vanilla, keep_trick), err_msg)
+
 
 class DeformConvTester(OpTester, unittest.TestCase):
     def expected_fn(self, x, weight, offset, mask, bias, stride=1, padding=0, dilation=1):
diff --git a/torchvision/ops/boxes.py b/torchvision/ops/boxes.py
index cfce618845a..d81fcda7e5c 100644
--- a/torchvision/ops/boxes.py
+++ b/torchvision/ops/boxes.py
@@ -1,6 +1,6 @@
 import torch
 from torch import Tensor
-from typing import Tuple
+from typing import List, Tuple
 from ._box_convert import _box_cxcywh_to_xyxy, _box_xyxy_to_cxcywh, _box_xywh_to_xyxy, _box_xyxy_to_xywh
 import torchvision
 from torchvision.extension import _assert_has_ops
@@ -62,18 +62,51 @@ def batched_nms(
             the elements that have been kept by NMS, sorted
             in decreasing order of scores
     """
-    if boxes.numel() == 0:
+    num_boxes = boxes.numel()
+
+    if num_boxes == 0:
         return torch.empty((0,), dtype=torch.int64, device=boxes.device)
-    # strategy: in order to perform NMS independently per class.
+    # Benchmarks that drove the following thresholds are at
+    # https://github.com/pytorch/vision/issues/1311#issuecomment-781329339
+    # Ideally for GPU we'd use a higher threshold
+    elif num_boxes > 4_000:
+        return _batched_nms_vanilla(boxes, scores, idxs, iou_threshold)
+    else:
+        return _batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold)
+
+
+def _batched_nms_coordinate_trick(
+    boxes: Tensor,
+    scores: Tensor,
+    idxs: Tensor,
+    iou_threshold: float,
+) -> Tensor:
+    # strategy: in order to perform NMS independently per class,
     # we add an offset to all the boxes. The offset is dependent
     # only on the class idx, and is large enough so that boxes
     # from different classes do not overlap
-    else:
-        max_coordinate = boxes.max()
-        offsets = idxs.to(boxes) * (max_coordinate + torch.tensor(1).to(boxes))
-        boxes_for_nms = boxes + offsets[:, None]
-        keep = nms(boxes_for_nms, scores, iou_threshold)
-        return keep
+    max_coordinate = boxes.max()
+    offsets = idxs.to(boxes) * (max_coordinate + torch.tensor(1).to(boxes))
+    boxes_for_nms = boxes + offsets[:, None]
+    keep = nms(boxes_for_nms, scores, iou_threshold)
+    return keep
+
+
+def _batched_nms_vanilla(
+    boxes: Tensor,
+    scores: Tensor,
+    idxs: Tensor,
+    iou_threshold: float,
+) -> Tensor:
+    # Based on Detectron2 implementation
+    result_mask = scores.new_zeros(scores.size(), dtype=torch.bool)
+    for id in torch.jit.annotate(List[int], torch.unique(idxs).cpu().tolist()):
+        mask = torch.where(idxs == id)[0].view(-1)
+        keep = nms(boxes[mask], scores[mask], iou_threshold)
+        result_mask[mask[keep]] = True
+    keep = torch.where(result_mask)[0].view(-1)
+    keep = keep[scores[keep].argsort(descending=True)]
+    return keep
 
 
 def remove_small_boxes(boxes: Tensor, min_size: float) -> Tensor:


Code review comments:
diff_hunk: @@ -62,18 +62,51 @@ def batched_nms(
             the elements that have been kept by NMS, sorted
             in decreasing order of scores
     """
-    if boxes.numel() == 0:
+    num_boxes = boxes.numel()
+
+    if num_boxes == 0:
         return torch.empty((0,), dtype=torch.int64, device=boxes.device)
-    # strategy: in order to perform NMS independently per class.
+    # Benchmarks that drove the following thresholds are at
+    # https://github.com/pytorch/vision/issues/1311#issuecomment-781329339
+    # Ideally for GPU we'd use a higher threshold
+    elif num_boxes > 4_000:
+        return _batched_nms_vanilla(boxes, scores, idxs, iou_threshold)
+    else:
+        return _batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold)
+
+
+def _batched_nms_coordinate_trick(
+    boxes: Tensor,
+    scores: Tensor,
+    idxs: Tensor,
+    iou_threshold: float,
+) -> Tensor:
+    # strategy: in order to perform NMS independently per class,
     # we add an offset to all the boxes. The offset is dependent
     # only on the class idx, and is large enough so that boxes
     # from different classes do not overlap
-    else:
-        max_coordinate = boxes.max()
-        offsets = idxs.to(boxes) * (max_coordinate + torch.tensor(1).to(boxes))
-        boxes_for_nms = boxes + offsets[:, None]
-        keep = nms(boxes_for_nms, scores, iou_threshold)
-        return keep
+    max_coordinate = boxes.max()
+    offsets = idxs.to(boxes) * (max_coordinate + torch.tensor(1).to(boxes))
+    boxes_for_nms = boxes + offsets[:, None]
+    keep = nms(boxes_for_nms, scores, iou_threshold)
+    return keep
+
+
+def _batched_nms_vanilla(
+    boxes: Tensor,
+    scores: Tensor,
+    idxs: Tensor,
+    iou_threshold: float,
+) -> Tensor:
+    # Based on Detectron2 implementation
+    result_mask = scores.new_zeros(scores.size(), dtype=torch.bool)
path: torchvision/ops/boxes.py
text: @user1:
should this be `torch.zeros_like(scores, dtype = torch.bool)`?