Rotated bounding box NMS implementation for CPU by zy1git · Pull Request #9450 · pytorch/vision

zy1git · 2026-03-23T04:28:43Z

Summary:
Implemented rotated box NMS (Non-Maximum Suppression) for CPU, adapted from Detectron2's nms_rotated implementation. The NMS algorithm is identical to standard NMS — sort by scores, suppress overlapping boxes — but uses single_box_iou_rotated for IoU computation instead of axis-aligned intersection. The public API follows the existing nms op pattern in TorchVision.

Test Plan:

Added TestNMSRotated test class adapted from Detectron2's test suite:

0° rotation test: rotated NMS with angle=0 should match reference horizontal NMS (IoU thresholds 0.2, 0.5, 0.8)
90° rotation test: rotated NMS with angle=90 and swapped width/height should match reference horizontal NMS
180° rotation test: rotated NMS with angle=180 should match reference horizontal NMS
TorchScript compatibility test

Results are compared using edit distance (≤ 1 allowed) to account for floating-point precision differences at IoU threshold boundaries.

Run pytest test/test_ops.py::TestNMSRotated -v
All tests pass locally.

pytorch-bot · 2026-03-23T04:28:47Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9450

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

viable/strict has been blocked for 5+ days

✅ No Failures

As of commit 54aee4c with merge base d7400a3 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

NicolasHug · 2026-03-23T10:54:53Z

torchvision/csrc/ops/cpu/nms_rotated_kernel.cpp

+
+      auto ovr = single_box_iou_rotated<scalar_t>(
+          dets[i].data_ptr<scalar_t>(), dets[j].data_ptr<scalar_t>());
+      if (ovr >= iou_threshold) {


Flagging that this is different from the iou threshold comparison we have in the non-rotated case:

vision/torchvision/csrc/ops/cpu/nms_kernel.cpp

Line 75 in d7400a3

if (ovr > iou_threshold) {

.

See my other comment about unifying the implementation, which should resolve this as a consequence.

Yes. It is resolved in the new commit after unifying the implementation.

NicolasHug · 2026-03-23T10:55:44Z

torchvision/csrc/ops/cpu/nms_rotated_kernel.cpp

+namespace {
+
+template <typename scalar_t>
+at::Tensor nms_rotated_cpu_kernel(


This is exactly the same implementation we already have for the non-rotated case, the only difference being the iou computation:

vision/torchvision/csrc/ops/cpu/nms_kernel.cpp

Line 10 in d7400a3

at::Tensor nms_kernel_impl(

Could we consider fusing the two implementations, perhaps templating over the iou computation function?

Yes, I had the same thought when I did the implementation. I wanted to stick to Detectron2's implementation in the first version to make sure it works correctly, then refactor. I have fused the two implementations in the new commit.

NicolasHug · 2026-03-23T10:56:46Z

torchvision/ops/boxes.py

    return torch.ops.torchvision.nms(boxes, scores, iou_threshold)


+def nms_rotated(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:


any reason to expose nms_rotated instead of just handling all this within a single nms function?

For iou, we chose not to expose iou_rotated at the Python layer.

Thanks for pointing this out. I fixed this in the new commit.

…gorithm is standard TorchVision code; attribution already in box_iou_rotated_utils.h

NicolasHug · 2026-03-24T16:32:01Z

torchvision/ops/boxes.py



-def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:
+def nms(boxes: Tensor, scores: Tensor, iou_threshold: float, fmt: str = "xyxy") -> Tensor:


Can we just infer the format base on the .shape attribute? If shape [-2] is 4 then it should be aligned, if 5 then it should be non-aligned?
I'm not sure, we'd have to verify that this actually doesn't break any other convension. Like, was shape[-2] == 5 even allowed before? If not, we're clear

This is a good point. I think it should be shape[-1] rather than shape[-2]. I checked shape[-1] == 5 was not allowed before (the kernel enforces size(1) == 4). And the format is unique for dimension 4 and 5, so this works. I have implemented it in the new commit.

NicolasHug · 2026-03-26T15:37:34Z

test/test_ops.py



+class TestNMSRotated:
+    def _reference_horizontal_nms(self, boxes, scores, iou_threshold):


Use TestNMS._reference_nms instead, you might have to make it a class method instead of an instance method

NicolasHug · 2026-03-26T15:40:12Z

test/test_ops.py

+            box_scores (N, 5): boxes in corner-form and probabilities.
+                (Note here 5 == 4 + 1, i.e., 4-dim horizontal box + 1-dim prob)


This will probably be resolved automatically when we use TestNMS._reference_nms, but flagging that this box_scores parameter doesn't exist.

NicolasHug · 2026-03-26T15:49:12Z

torchvision/ops/boxes.py

+    else:
+        raise ValueError(
+            f"boxes should have 4 (axis-aligned) or 5 (rotated) elements in the last dimension, got {boxes.size(-1)}"
+        )


let's also make sure to address and test batched_nms just below

NicolasHug · 2026-03-26T15:49:48Z

test/test_ops.py

+        m, n = len(keep1), len(keep2)
+
+        # edit distance with DP
+        f = [np.arange(n + 1), np.arange(n + 1)]


don't use numpy

NicolasHug · 2026-03-26T15:51:55Z

test/test_ops.py

+    @pytest.mark.parametrize("iou", (0.2, 0.5, 0.8))
+    def test_nms_rotated_0_degree(self, iou):
+        N = 1000
+        boxes, scores = self._create_tensors(N)


Here and in other tests, call torch.manual_seed(0) so that the test is not dependent on RNG.

NicolasHug · 2026-03-26T15:52:24Z

test/test_ops.py

+
+        keep_ref = self._reference_horizontal_nms(boxes, scores, iou)
+        keep = ops.nms(rotated_boxes, scores, iou)
+        assert self._nms_edit_distance(keep, keep_ref) <= 1


Locally, I am able to use torch.assert_close(..., atol=0) on many random seeds, we should consider using this instead of edit distance.

NicolasHug · 2026-03-26T15:54:44Z

test/test_ops.py

+        return boxes, scores
+
+    @pytest.mark.parametrize("iou", (0.2, 0.5, 0.8))
+    def test_nms_rotated_0_degree(self, iou):


This is comparing our rotated implementation against _reference_horizontal_nms. We should also have a test that uses our non-rotated nms implementation as the reference.

NicolasHug · 2026-03-26T15:55:52Z

test/test_ops.py

+        return torch.as_tensor(picked)
+
+    @staticmethod
+    def _nms_edit_distance(keep1, keep2):


Remove if we don't need it based on https://github.com/pytorch/vision/pull/9450/changes#r2995974410

NicolasHug · 2026-03-26T15:56:25Z

torchvision/csrc/ops/cpu/nms_kernel.cpp

-      auto inter = w * h;
-      auto ovr = inter / (iarea + areas[j] - inter);
+
+      auto ovr = iou_func.compare(j);


Nit: call this compute rather than compare

NicolasHug · 2026-03-26T15:57:36Z

torchvision/csrc/ops/cpu/nms_kernel.cpp

+
+  scalar_t ix1, iy1, ix2, iy2, iarea;
+
+  AABBIoU(const at::Tensor& dets) {


Call this NonRotatedIoU or AlignedIoU

NicolasHug · 2026-03-26T16:00:30Z

torchvision/csrc/ops/cpu/nms_kernel.cpp

+
+  RotatedIoU(const at::Tensor& dets) : dets_ptr(&dets) {}
+
+  int64_t cached_i;


just call this i

NMS implementation for CPU

bfb8321

zy1git marked this pull request as draft March 23, 2026 04:28

meta-cla bot added the cla signed label Mar 23, 2026

Zhitao Yu added 2 commits March 22, 2026 22:58

add tests

8875875

Merge remote-tracking branch 'origin/main' into rotated-NMS

ae5fb41

NicolasHug reviewed Mar 23, 2026

View reviewed changes

Zhitao Yu added 6 commits March 24, 2026 01:17

fuse two implementations

0e924ed

Remove Detectron2 license header from nms_kernel.cpp since the NMS al…

667e1b4

…gorithm is standard TorchVision code; attribution already in box_iou_rotated_utils.h

revert the unnecessary linting

f01ca8b

Preserve original IoU computation in fused NMS implementation

40adc7f

Add box caching in fused NMS to match original memory access pattern

eb4aefe

fix the test failure

5499015

NicolasHug reviewed Mar 24, 2026

View reviewed changes

Zhitao Yu added 3 commits March 24, 2026 19:44

remove the torchscript test and err_msg

c940345

Remove the fmt parameter from nms function

17a54cc

fix the test failures

54aee4c

zy1git marked this pull request as ready for review March 26, 2026 04:15

NicolasHug reviewed Mar 26, 2026

View reviewed changes

test/test_ops.py

m, n = len(keep1), len(keep2)

# edit distance with DP

f = [np.arange(n + 1), np.arange(n + 1)]

Copy link

Member

NicolasHug Mar 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

don't use numpy

NicolasHug reviewed Mar 26, 2026

View reviewed changes

		return torch.ops.torchvision.nms(boxes, scores, iou_threshold)


		def nms_rotated(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:



		def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:
		def nms(boxes: Tensor, scores: Tensor, iou_threshold: float, fmt: str = "xyxy") -> Tensor:



		class TestNMSRotated:
		def _reference_horizontal_nms(self, boxes, scores, iou_threshold):

		box_scores (N, 5): boxes in corner-form and probabilities.
		(Note here 5 == 4 + 1, i.e., 4-dim horizontal box + 1-dim prob)


		scalar_t ix1, iy1, ix2, iy2, iarea;

		AABBIoU(const at::Tensor& dets) {


		RotatedIoU(const at::Tensor& dets) : dets_ptr(&dets) {}

		int64_t cached_i;

Conversation

zy1git commented Mar 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Mar 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9450

❗ 1 Active SEVs

✅ No Failures

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

zy1git commented Mar 23, 2026 •

edited

Loading

pytorch-bot bot commented Mar 23, 2026 •

edited

Loading