Removes cls_outputs and box_outputs from postprocessing results.

PiperOrigin-RevId: 276586434

official/vision/detection/executor/detection_executor.py

@@ -95,9 +95,9 @@ class DetectionDistributedExecutor(executor.DistributedExecutor):
         inputs, labels = inputs
         model_outputs = model(inputs, training=False)
         if self._predict_post_process_fn:
-          labels, model_outputs = self._predict_post_process_fn(
+          labels, prediction_outputs = self._predict_post_process_fn(
               labels, model_outputs)
-        return labels, model_outputs
+        return labels, prediction_outputs
 
       labels, outputs = strategy.experimental_run_v2(
           _test_step_fn, args=(next(iterator),))

official/vision/detection/modeling/retinanet_model.py

@@ -40,10 +40,8 @@ class COCOMetrics(object):
     self._evaluator = eval_factory.evaluator_generator(params.eval)
 
   def update_state(self, y_true, y_pred):
-    labels, outputs = y_true, y_pred
-
-    labels = tf.nest.map_structure(lambda x: x.numpy(), labels)
-    outputs = tf.nest.map_structure(lambda x: x.numpy(), outputs)
+    labels = tf.nest.map_structure(lambda x: x.numpy(), y_true)
+    outputs = tf.nest.map_structure(lambda x: x.numpy(), y_pred)
     groundtruths = {}
     predictions = {}
     for key, val in outputs.items():
@@ -161,14 +159,16 @@ class RetinanetModel(base_model.Model):
     boxes, scores, classes, valid_detections = self._generate_detections_fn(
         inputs=(outputs['box_outputs'], outputs['cls_outputs'],
                 labels['anchor_boxes'], labels['image_info'][:, 1:2, :]))
-    outputs.update({
+    # Discards the old output tensors to save memory. The `cls_outputs` and
+    # `box_outputs` are pretty big and could potentiall lead to memory issue.
+    outputs = {
         'source_id': labels['groundtruths']['source_id'],
         'image_info': labels['image_info'],
         'num_detections': valid_detections,
         'detection_boxes': boxes,
         'detection_classes': classes,
         'detection_scores': scores,
-    })
+    }
 
     if 'groundtruths' in labels:
       labels['source_id'] = labels['groundtruths']['source_id']