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

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.

❌ 2 New Failures

As of commit aea073c with merge base d7400a3 ():

NEW FAILURES - The following jobs have failed:

Build Windows Wheels / pytorch/vision / build-wheel-py3_10-cuda13_2 (gh)
Build Windows Wheels / pytorch/vision / upload / upload-wheel-py3_10-cuda13_2 (gh)
Unable to download artifact(s): Artifact not found for name: pytorch_vision__3.10_cu132_x64

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.

Good point. box_scores doesn't exist as a parameter. The docstring is incorrect. Detectron2 has the same issue: https://github.com/facebookresearch/detectron2/blob/main/tests/layers/test_nms_rotated.py#L46-L48

torchvision/ops/boxes.py

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

Got it. Will pay attention to this in future.

test/test_ops.py

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.

Fixed in the new commit.

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.

I agree. Fixed in the new commit.

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

Removed in the new commit.

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

fixed in the new commit.

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

fixed in the new commit.

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

fixed in the new commit.

….cpp

NicolasHug · 2026-03-30T15:56:47Z

test/test_ops.py

        torch.testing.assert_close(iou_cpu, iou_cuda.cpu(), atol=1e-5, rtol=1e-5)


+class TestNMSRotated:


Now that we are exposing this through the same nms() method, we should just move these tests within TestNMS.

Addressed in the new commit.

NicolasHug · 2026-03-30T15:57:20Z

test/test_ops.py

 class TestNMS:
-    def _reference_nms(self, boxes, scores, iou_threshold):
+    @classmethod
+    def _reference_nms(cls, boxes, scores, iou_threshold):


now that we're moving the tests, we should rename this to _reference_aligned_nms or something like that

Addressed in the new commit.

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

test/test_ops.py

+        boxes, scores = self._create_tensors(N)
+        rotated_boxes = torch.zeros(N, 5)
+        rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
+        rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
+        # Swap width and height for 90 degrees so reference horizontal NMS can be used
+        rotated_boxes[:, 2] = boxes[:, 3] - boxes[:, 1]
+        rotated_boxes[:, 3] = boxes[:, 2] - boxes[:, 0]
+        rotated_boxes[:, 4] = 90


Instead of duplicating this logic everywhere, let's create aligned boxes, call a box format conversion function, and reset the column angle as desired. We probably want to do that in the _create_tensor() helper which hsould be renamed to _create_rotated_boxes()

NicolasHug · 2026-03-30T16:03:25Z

test/test_ops.py

+        torch.testing.assert_close(keep, keep_non_rotated, atol=0, rtol=0)
+
+    @pytest.mark.parametrize("iou", (0.2, 0.5, 0.8))
+    def test_nms_rotated_90_degrees(self, iou):


parametrize over the angle instead of duplicating the test N times. Also, it doesn't seem necessary for the angles to be a multiple of 90 for this to work, I think it'll work with any angle value, the reason everything works is because the angle is the same across boxes, but its actual value doesn't matter (I could be wrong, please double check)

I parametrized over the angle in the new commit. However, I think we can only parametrize for 0, 90, and 180. For other angles like 45°, even though each box is rotated by the same angle, they rotate around their own centers (which are at different positions), so the overlapped area changes and we can't compare against axis-aligned NMS. For 90, we swap width and height so the rotated boxes cover the same region as the axis-aligned boxes, so the NMS are the same.

NicolasHug · 2026-03-30T16:05:20Z

test/test_ops.py

+        torch.testing.assert_close(keep, keep_ref, atol=0, rtol=0)
+        keep_non_rotated = ops.nms(boxes, scores, iou)
+        torch.testing.assert_close(keep, keep_non_rotated, atol=0, rtol=0)
+


we should add tests for when the boxes have a different rotation angle, currently all tests have the same angle.

I added test_nms_rotated_different_angles, which verifies basic NMS properties since we can't compare against axis-aligned NMS when boxes have different angles. And also I added test_nms_rotated_specific_angles to show that the NMS worked as expected for different specific angles. Please feel free to let me know what extra tests for different rotation angles you want me to add.

NicolasHug · 2026-03-30T16:12:28Z

test/test_ops.py

+        torch.testing.assert_close(keep, keep_non_rotated, atol=0, rtol=0)
+
+    @pytest.mark.parametrize("iou", (0.2, 0.5, 0.8))
+    def test_batched_nms_rotated_0_degree(self, iou):


Let's also parametrize the existing test_batched_nms_implementations test over the format (one rotated, one non-rotated).

Addressed in the new commit

…different rotation angles per box

…he format

NicolasHug · 2026-04-02T15:36:57Z

test/test_ops.py

+        backup = rotated_boxes.clone()
+        keep_non_rotated = ops.batched_nms(boxes, scores, idxs, iou)
+        keep = ops.batched_nms(rotated_boxes, scores, idxs, iou)
+        assert torch.allclose(rotated_boxes, backup)


torch.testing.assert_close

fixed in the new commit.

NicolasHug · 2026-04-02T15:37:17Z

test/test_ops.py

+    def test_nms_rotated_different_angles(self, iou):
+        torch.manual_seed(0)
+        N = 1000
+        boxes, rotated_boxes, scores = self._create_rotated_boxes(N)


boxes are not used

fixed by "_" in the new commit.

NicolasHug · 2026-04-02T15:41:09Z

test/test_ops.py

+        torch.testing.assert_close(keep, keep_non_rotated, atol=0, rtol=0)
+
+    @pytest.mark.parametrize("iou", (0.2, 0.5, 0.8))
+    def test_batched_nms_rotated_0_degree(self, iou):


sounds like this could be tested on 90 and 180 as well. Consider parametrizing with test_nms_rotated if it doesn't complicate the code too much.

Addressed in the new commit.

NicolasHug · 2026-04-02T15:42:07Z

test/test_ops.py

+        if angle == 90:
+            rotated_boxes[:, 2] = cxcywh[:, 3]
+            rotated_boxes[:, 3] = cxcywh[:, 2]
+        else:
+            rotated_boxes[:, 2] = cxcywh[:, 2]
+            rotated_boxes[:, 3] = cxcywh[:, 3]


This trick should be explained. I suspect we'll want to move it out of the _create_rotated_boxes as it's mainly used to make the existing test_nms_rotated work.

Addressed in the new commit.

NicolasHug · 2026-04-02T15:45:06Z

torchvision/csrc/ops/cpu/nms_kernel.cpp

+        (*dets_ptr)[i].template data_ptr<scalar_t>(),
+        (*dets_ptr)[j].template data_ptr<scalar_t>());


can you explain the .template syntax?

I added .template to fix a CI failure we hit on the macOS build. It explicitly indicates that data_ptr<scalar_t> is a template method. Without it, the compiler can't tell whether data_ptr<scalar_t> is a template method call or a less-than comparison.

…ve boxes in test_nms_rotated_different_angles function

…_boxes

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

torchvision/ops/boxes.py Show resolved Hide resolved

NicolasHug reviewed Mar 26, 2026

View reviewed changes

test/test_ops.py Outdated Show resolved Hide resolved

NicolasHug reviewed Mar 26, 2026

View reviewed changes

Reuse TestNMS._reference_nms in TestNMSRotated

c5cf1c4

Zhitao Yu added 4 commits March 27, 2026 01:37

address the comment

9c96e26

change the variable name

4a993c8

address more comments on the file torchvision/csrc/ops/cpu/nms_kernel…

84e960c

….cpp

add batched_nms and the test

9ef0c0b

NicolasHug reviewed Mar 31, 2026

View reviewed changes

Zhitao Yu added 5 commits March 31, 2026 10:56

Rename _reference_nms to _reference_aligned_nms

4da74a8

Merge TestNMSRotated into TestNMS

2ea193b

combine tests, refactor duplicated parts and add a test for NMS with …

c249d35

…different rotation angles per box

add test_nms_rotated_specific_angles

50afe5c

parametrize the existing test_batched_nms_implementations test over t…

8259471

…he format

NicolasHug reviewed Apr 2, 2026

View reviewed changes

Zhitao Yu added 3 commits April 2, 2026 21:50

address the comments by using the torch.testing.assert_close and remo…

bdd502d

…ve boxes in test_nms_rotated_different_angles function

parametrizing the angles for the test_batched_nms_rotated function

7f5d72b

add explanation for angle = 90 and move it out of the _create_rotated…

aea073c

…_boxes

		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;

		torch.testing.assert_close(iou_cpu, iou_cuda.cpu(), atol=1e-5, rtol=1e-5)


		class TestNMSRotated:

		(*dets_ptr)[i].template data_ptr<scalar_t>(),
		(*dets_ptr)[j].template data_ptr<scalar_t>());

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

❌ 2 New 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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

zy1git commented Mar 23, 2026 •

edited

Loading

pytorch-bot bot commented Mar 23, 2026 •

edited

Loading

zy1git Apr 1, 2026 •

edited

Loading