class_docstr: Test DETRModel with auxilary loss.
class_docstr_tok: ["Test","DETRModel","with","auxilary","loss","."]
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):
      """Test loss_function by checking its output's dictionary format."""
  
      # test loss function in the pmapped setup
      loss_function_pmapped = jax.pmap(
          self.model.loss_function, axis_name="batch"
      )
  
      matches = jax_utils.replicate(
          # fake matching for the  final output +  2 aux outputs
          [self.indices]
          * 3
      )
      outputs_replicated, batch_replicated = (
          jax_utils.replicate(self.outputs),
          jax_utils.replicate(self.batch),
      )
      total_loss, metrics_dict = loss_function_pmapped(
          outputs_replicated, batch_replicated, matches=matches
      )
  
      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 self.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 going to be included in the total object dtection loss
          if (
              "_unscaled" not in k
              and b in self.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=3
      )
  
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  from flax import jax_utils
  from flax import traverse_util
  from flax.core import unfreeze
  import jax
  from jax import random
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.baselines.detr import model
label_desc: Use image node_id 2 for calling the DETRModelTestWithAuxLoss obj's underlying member method code with example usage: obj.test_loss_function() without return types