Support all NMS versions besides tflite for Keras QAT RetinaNet model.

PiperOrigin-RevId: 519777521

Author: fyangf

official/projects/qat/vision/modeling/factory.py

@@ -201,7 +201,9 @@ def build_qat_retinanet(
   if quantization.quantize_detection_head:
     # Call the model with dummy input to build the head part.
     dummpy_input = tf.zeros([1] + model_config.input_size)
-    optimized_model(dummpy_input, training=True)
+    height, width, _ = model_config.input_size
+    image_shape = [[height, width]]
+    optimized_model.call(dummpy_input, image_shape=image_shape, training=False)
     helper.copy_original_weights(model.head, optimized_model.head)
   return optimized_model
 

official/projects/qat/vision/serving/export_module.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 """Export modules for QAT model serving/inference."""
-from absl import logging
 import tensorflow as tf
 
 from official.projects.qat.vision.modeling import factory as qat_factory
@@ -50,11 +49,6 @@ class DetectionModule(detection.DetectionModule):
   """Detection Module."""
 
   def _build_model(self):
-    if self.params.task.model.detection_generator.nms_version != 'tflite':
-      self.params.task.model.detection_generator.nms_version = 'tflite'
-      logging.info('Set `nms_version` to `tflite` because only TFLite NMS is '
-                   'supported for QAT detection models.')
-
     model = super()._build_model()
 
     if isinstance(self.params.task.model, configs.retinanet.RetinaNet):

official/vision/modeling/layers/detection_generator.py

@@ -1310,7 +1310,10 @@ def _decode_multilevel_outputs_and_pre_nms_top_k(
     levels = list(raw_boxes.keys())
     min_level = int(min(levels))
     max_level = int(max(levels))
-    clip_shape = tf.expand_dims(tf.expand_dims(image_shape, axis=1), axis=1)
+    if image_shape is not None:
+      clip_shape = tf.expand_dims(tf.expand_dims(image_shape, axis=1), axis=1)
+    else:
+      clip_shape = None
     for i in range(max_level, min_level - 1, -1):
       (
           batch_size,
@@ -1330,13 +1333,15 @@ def _decode_multilevel_outputs_and_pre_nms_top_k(
           unsharded_w * num_anchors_per_locations,
           4,
       ]
-      decoded_boxes = box_ops.clip_boxes(
-          box_ops.decode_boxes(
-              tf.reshape(raw_boxes[str(i)], boxes_shape),
-              tf.reshape(anchor_boxes[str(i)], boxes_shape),
-          ),
-          clip_shape,
+      decoded_boxes = box_ops.decode_boxes(
+          tf.reshape(raw_boxes[str(i)], boxes_shape),
+          tf.reshape(anchor_boxes[str(i)], boxes_shape),
       )
+      if clip_shape is not None:
+        decoded_boxes = box_ops.clip_boxes(
+            decoded_boxes,
+            clip_shape,
+        )
       for raw_scores_i, decoded_boxes_i in edgetpu.shard_tensors(
           1, block, (raw_scores[str(i)], decoded_boxes)
       ):