class_docstr: Test DeformableDETRModel loss test.
class_docstr_tok: ["Test","DeformableDETRModel","loss","test","."]
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_and_metrics_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.
      expected_metrics_keys = [
          "loss_class",
          "loss_bbox",
          "loss_giou",
          "total_loss",
      ]
      for i in range(self.num_decoder_layers - 1):
          expected_metrics_keys += [
              f"loss_class_aux{i}",
              f"loss_bbox_aux{i}",
              f"loss_giou_aux{i}",
          ]
      self.assertSameElements(
          expected_metrics_keys, metrics_dict.keys()
      )
  
      # Since 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 loss.
          if 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
      )
  
file_path: scenic/scenic/projects/baselines/deformable_detr/tests/test_model.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  from flax import jax_utils
  import jax
  from jax import random
  import jax.numpy
  import numpy
  from scenic.projects.baselines.deformable_detr.configs import mini_config
  from scenic.projects.baselines.deformable_detr.model import compute_cost
  from scenic.projects.baselines.deformable_detr.model import DeformableDETRModel
label_desc: Use image node_id 2 for calling the DeformableDETRModelLossTest obj's underlying member method code with example usage: obj.test_loss_function() without return types