Add more tests to NMS (#2279)

* Add more tests to NMS * Fix lint
pytorch · Jun 1, 2020 · 34810c0 · 34810c0
1 parent b40f49f
commit 34810c0
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/test/test_ops.py b/test/test_ops.py
@@ -393,6 +393,10 @@ def test_nms(self):
             keep_ref = self.reference_nms(boxes, scores, iou)
             keep = ops.nms(boxes, scores, iou)
             self.assertTrue(torch.allclose(keep, keep_ref), err_msg.format(iou))
+        self.assertRaises(RuntimeError, ops.nms, torch.rand(4), torch.rand(3), 0.5)
+        self.assertRaises(RuntimeError, ops.nms, torch.rand(3, 5), torch.rand(3), 0.5)
+        self.assertRaises(RuntimeError, ops.nms, torch.rand(3, 4), torch.rand(3, 2), 0.5)
+        self.assertRaises(RuntimeError, ops.nms, torch.rand(3, 4), torch.rand(4), 0.5)
 
     @unittest.skipIf(not torch.cuda.is_available(), "CUDA unavailable")
     def test_nms_cuda(self):

diff --git a/torchvision/csrc/nms.h b/torchvision/csrc/nms.h
@@ -12,6 +12,24 @@ at::Tensor nms(
     const at::Tensor& dets,
     const at::Tensor& scores,
     const double iou_threshold) {
+  TORCH_CHECK(
+      dets.dim() == 2, "boxes should be a 2d tensor, got ", dets.dim(), "D");
+  TORCH_CHECK(
+      dets.size(1) == 4,
+      "boxes should have 4 elements in dimension 1, got ",
+      dets.size(1));
+  TORCH_CHECK(
+      scores.dim() == 1,
+      "scores should be a 1d tensor, got ",
+      scores.dim(),
+      "D");
+  TORCH_CHECK(
+      dets.size(0) == scores.size(0),
+      "boxes and scores should have same number of elements in ",
+      "dimension 0, got ",
+      dets.size(0),
+      " and ",
+      scores.size(0));
   if (dets.is_cuda()) {
 #if defined(WITH_CUDA)
     if (dets.numel() == 0) {