Internal change

PiperOrigin-RevId: 394750988

official/vision/beta/configs/maskrcnn.py

@@ -132,6 +132,7 @@ class DetectionGenerator(hyperparams.Config):
   nms_iou_threshold: float = 0.5
   max_num_detections: int = 100
   use_batched_nms: bool = False
+  use_cpu_nms: bool = False
 
 
 @dataclasses.dataclass

official/vision/beta/configs/retinanet.py

@@ -113,6 +113,7 @@ class DetectionGenerator(hyperparams.Config):
   nms_iou_threshold: float = 0.5
   max_num_detections: int = 100
   use_batched_nms: bool = False
+  use_cpu_nms: bool = False
 
 
 @dataclasses.dataclass

official/vision/beta/modeling/factory.py

@@ -196,7 +196,8 @@ def build_maskrcnn(
       pre_nms_score_threshold=generator_config.pre_nms_score_threshold,
       nms_iou_threshold=generator_config.nms_iou_threshold,
       max_num_detections=generator_config.max_num_detections,
-      use_batched_nms=generator_config.use_batched_nms)
+      use_batched_nms=generator_config.use_batched_nms,
+      use_cpu_nms=generator_config.use_cpu_nms)
 
   if model_config.include_mask:
     mask_head = instance_heads.MaskHead(
@@ -293,7 +294,8 @@ def build_retinanet(
       pre_nms_score_threshold=generator_config.pre_nms_score_threshold,
       nms_iou_threshold=generator_config.nms_iou_threshold,
       max_num_detections=generator_config.max_num_detections,
-      use_batched_nms=generator_config.use_batched_nms)
+      use_batched_nms=generator_config.use_batched_nms,
+      use_cpu_nms=generator_config.use_cpu_nms)
 
   model = retinanet_model.RetinaNetModel(
       backbone,

official/vision/beta/modeling/layers/detection_generator.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Contains definitions of generators to generate the final detections."""
+import contextlib
 from typing import List, Optional, Mapping
 # Import libraries
 import tensorflow as tf
@@ -404,6 +405,7 @@ class DetectionGenerator(tf.keras.layers.Layer):
                nms_iou_threshold: float = 0.5,
                max_num_detections: int = 100,
                use_batched_nms: bool = False,
+               use_cpu_nms: bool = False,
                **kwargs):
     """Initializes a detection generator.
 
@@ -420,6 +422,7 @@ class DetectionGenerator(tf.keras.layers.Layer):
         generate.
       use_batched_nms: A `bool` of whether or not use
         `tf.image.combined_non_max_suppression`.
+      use_cpu_nms: A `bool` of whether or not enforce NMS to run on CPU.
       **kwargs: Additional keyword arguments passed to Layer.
     """
     self._config_dict = {
@@ -429,6 +432,7 @@ class DetectionGenerator(tf.keras.layers.Layer):
         'nms_iou_threshold': nms_iou_threshold,
         'max_num_detections': max_num_detections,
         'use_batched_nms': use_batched_nms,
+        'use_cpu_nms': use_cpu_nms,
     }
     super(DetectionGenerator, self).__init__(**kwargs)
 
@@ -513,6 +517,13 @@ class DetectionGenerator(tf.keras.layers.Layer):
           'decoded_box_scores': box_scores,
       }
 
+    # Optionally force the NMS be run on CPU.
+    if self._config_dict['use_cpu_nms']:
+      nms_context = tf.device('cpu:0')
+    else:
+      nms_context = contextlib.nullcontext()
+
+    with nms_context:
       if self._config_dict['use_batched_nms']:
         (nmsed_boxes, nmsed_scores, nmsed_classes, valid_detections) = (
             _generate_detections_batched(
@@ -560,6 +571,7 @@ class MultilevelDetectionGenerator(tf.keras.layers.Layer):
                nms_iou_threshold: float = 0.5,
                max_num_detections: int = 100,
                use_batched_nms: bool = False,
+               use_cpu_nms: bool = False,
                **kwargs):
     """Initializes a multi-level detection generator.
 
@@ -576,6 +588,7 @@ class MultilevelDetectionGenerator(tf.keras.layers.Layer):
         generate.
       use_batched_nms: A `bool` of whether or not use
         `tf.image.combined_non_max_suppression`.
+      use_cpu_nms: A `bool` of whether or not enforce NMS to run on CPU.
       **kwargs: Additional keyword arguments passed to Layer.
     """
     self._config_dict = {
@@ -585,6 +598,7 @@ class MultilevelDetectionGenerator(tf.keras.layers.Layer):
         'nms_iou_threshold': nms_iou_threshold,
         'max_num_detections': max_num_detections,
         'use_batched_nms': use_batched_nms,
+        'use_cpu_nms': use_cpu_nms,
     }
     super(MultilevelDetectionGenerator, self).__init__(**kwargs)
 
@@ -710,9 +724,17 @@ class MultilevelDetectionGenerator(tf.keras.layers.Layer):
           'decoded_box_attributes': attributes,
       }
 
+    # Optionally force the NMS to run on CPU.
+    if self._config_dict['use_cpu_nms']:
+      nms_context = tf.device('cpu:0')
+    else:
+      nms_context = contextlib.nullcontext()
+
+    with nms_context:
       if self._config_dict['use_batched_nms']:
         if raw_attributes:
-        raise ValueError('Attribute learning is not supported for batched NMS.')
+          raise ValueError(
+              'Attribute learning is not supported for batched NMS.')
 
         (nmsed_boxes, nmsed_scores, nmsed_classes, valid_detections) = (
             _generate_detections_batched(

official/vision/beta/modeling/layers/detection_generator_test.py

@@ -43,11 +43,9 @@ class SelectTopKScoresTest(tf.test.TestCase):
 class DetectionGeneratorTest(
     parameterized.TestCase, tf.test.TestCase):
 
-  @parameterized.parameters(
-      (True),
-      (False),
-  )
-  def testDetectionsOutputShape(self, use_batched_nms):
+  @parameterized.product(
+      use_batched_nms=[True, False], use_cpu_nms=[True, False])
+  def testDetectionsOutputShape(self, use_batched_nms, use_cpu_nms):
     max_num_detections = 100
     num_classes = 4
     pre_nms_top_k = 5000
@@ -60,6 +58,7 @@ class DetectionGeneratorTest(
         'nms_iou_threshold': 0.5,
         'max_num_detections': max_num_detections,
         'use_batched_nms': use_batched_nms,
+        'use_cpu_nms': use_cpu_nms,
     }
     generator = detection_generator.DetectionGenerator(**kwargs)
 
@@ -99,6 +98,7 @@ class DetectionGeneratorTest(
         'nms_iou_threshold': 0.5,
         'max_num_detections': 10,
         'use_batched_nms': False,
+        'use_cpu_nms': False,
     }
     generator = detection_generator.DetectionGenerator(**kwargs)
 
@@ -116,16 +116,20 @@ class MultilevelDetectionGeneratorTest(
     parameterized.TestCase, tf.test.TestCase):
 
   @parameterized.parameters(
-      (True, False),
-      (False, False),
-      (False, True),
+      (True, False, True),
+      (True, False, False),
+      (False, False, True),
+      (False, False, False),
+      (False, True, True),
+      (False, True, False),
   )
-  def testDetectionsOutputShape(self, use_batched_nms, has_att_heads):
+  def testDetectionsOutputShape(self, use_batched_nms, has_att_heads,
+                                use_cpu_nms):
     min_level = 4
     max_level = 6
     num_scales = 2
     max_num_detections = 100
-    aspect_ratios = [1.0, 2.0,]
+    aspect_ratios = [1.0, 2.0]
     anchor_scale = 2.0
     output_size = [64, 64]
     num_classes = 4
@@ -139,6 +143,7 @@ class MultilevelDetectionGeneratorTest(
         'nms_iou_threshold': 0.5,
         'max_num_detections': max_num_detections,
         'use_batched_nms': use_batched_nms,
+        'use_cpu_nms': use_cpu_nms,
     }
 
     input_anchor = anchor.build_anchor_generator(min_level, max_level,
@@ -219,6 +224,7 @@ class MultilevelDetectionGeneratorTest(
         'nms_iou_threshold': 0.5,
         'max_num_detections': 10,
         'use_batched_nms': False,
+        'use_cpu_nms': False,
     }
     generator = detection_generator.MultilevelDetectionGenerator(**kwargs)