class_docstr: Test ObjectDetectionWithMatchingModel.
class_docstr_tok: ["Test","ObjectDetectionWithMatchingModel","."]
docstr: Test loss_function by checking its output's dictionary format.
docstr_tok: ["Test","loss_function","by","checking","its","output","'s","dictionary","format","."]
code: def test_loss_function(self, num_boxes):
      """Test loss_function by checking its output's dictionary format."""
      model = MyObjectDetectionWithMatchingModel()
      outputs, batch, indices = fake_model_outputs_batch(num_boxes)
      outputs_noaux = dict(outputs)
      outputs_noaux.pop("aux_outputs")
  
      outputs_replicated, batch_replicated = (
          jax_utils.replicate(outputs_noaux),
          jax_utils.replicate(batch),
      )
  
      # Test loss function in the pmapped setup:
      loss_function_pmapped = jax.pmap(
          model.loss_function, axis_name="batch"
      )
      outputs_replicated = jax_utils.replicate(outputs)
  
      indices_replicated = jax_utils.replicate(
          # Fake matching for the  final output +  2 aux outputs:
          [indices]
          * 3
      )
      total_loss, metrics_dict = loss_function_pmapped(
          outputs_replicated, batch_replicated, indices_replicated
      )
  
      total_loss, metrics_dict = (
          jax_utils.unreplicate(total_loss),
          jax_utils.unreplicate(metrics_dict),
      )
  
      # Collect what keys we expect to find in the metrics_dict:
      base = [
          "class_accuracy",
          "loss_class",
          "loss_bbox",
          "loss_giou",
          "total_loss",
          "loss_class_aux_0",
          "loss_bbox_aux_0",
          "loss_giou_aux_0",
          "loss_class_aux_1",
          "loss_bbox_aux_1",
          "loss_giou_aux_1",
          "class_accuracy_not_pad",
      ]
      base_unscaled = []
      for b in base:
          if b.split("_aux_")[0] in model.loss_terms_weights.keys():
              base_unscaled.append(b + "_unscaled")
          else:
              base_unscaled.append(b)
      base_scaled = [
          "loss_class",
          "loss_bbox",
          "loss_giou",
          "loss_class_aux_0",
          "loss_bbox_aux_0",
          "loss_giou_aux_0",
          "loss_class_aux_1",
          "loss_bbox_aux_1",
          "loss_giou_aux_1",
      ]
      expected_metrics_keys = base_unscaled + base_scaled
      self.assertSameElements(
          expected_metrics_keys, metrics_dict.keys()
      )
  
      # Because weight decay is not used, the following must hold:
      object_detection_loss = 0
      for k in metrics_dict.keys():
          b = k.split("_aux_")[0]
          # If this loss is included in the total object detection loss...
          if (
              "_unscaled" not in k
              and b in model.loss_terms_weights.keys()
          ):
              # ...get the normalizer for this loss:
              object_detection_loss += (
                  # Already scaled loss term / loss term normalizer:
                  metrics_dict[k][0]
                  / metrics_dict[k][1]
              )
      self.assertAlmostEqual(
          total_loss, object_detection_loss, places=5
      )
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from flax import jax_utils
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.baselines.detr import detr_base_model
label_desc: Use image node_id 5 for calling the TestObjectDetectionWithMatchingModel obj's underlying member method code with example usage: obj.test_loss_function(num_boxes) without return types