From ddb6618bd81eb792d9c7d5a19871977209874ef7 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Thu, 20 Feb 2025 09:20:30 -0800 Subject: [PATCH 1/2] Update threshold for batched NMS --- torchvision/ops/boxes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/ops/boxes.py b/torchvision/ops/boxes.py index 9674d5bfa1d..ec6934618ab 100644 --- a/torchvision/ops/boxes.py +++ b/torchvision/ops/boxes.py @@ -78,7 +78,7 @@ def batched_nms( _log_api_usage_once(batched_nms) # Benchmarks that drove the following thresholds are at # https://github.com/pytorch/vision/issues/1311#issuecomment-781329339 - if boxes.numel() > (4000 if boxes.device.type == "cpu" else 20000) and not torchvision._is_tracing(): + if boxes.numel() > (4000 if boxes.device.type == "cpu" else 100_000) and not torchvision._is_tracing(): return _batched_nms_vanilla(boxes, scores, idxs, iou_threshold) else: return _batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold) From f3b94cf20123d041d4ab3bd1d8ec69937c61c69d Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Thu, 20 Feb 2025 09:21:09 -0800 Subject: [PATCH 2/2] Add comment --- torchvision/ops/boxes.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torchvision/ops/boxes.py b/torchvision/ops/boxes.py index ec6934618ab..48df4d85cc7 100644 --- a/torchvision/ops/boxes.py +++ b/torchvision/ops/boxes.py @@ -78,6 +78,7 @@ def batched_nms( _log_api_usage_once(batched_nms) # Benchmarks that drove the following thresholds are at # https://github.com/pytorch/vision/issues/1311#issuecomment-781329339 + # and https://github.com/pytorch/vision/pull/8925 if boxes.numel() > (4000 if boxes.device.type == "cpu" else 100_000) and not torchvision._is_tracing(): return _batched_nms_vanilla(boxes, scores, idxs, iou_threshold) else: