class_docstr: Tests the DeformableDETR trainer on single device setup.
class_docstr_tok: ["Tests","the","DeformableDETR","trainer","on","single","device","setup","."]
docstr: Test training for a few steps of a mini DETR on a mini COCO dataset.
docstr_tok: ["Test","training","for","a","few","steps","of","a","mini","DETR","on","a","mini","COCO","dataset","."]
code: def test_deformable_detr_trainer(self):
      """Test training for a few steps of a mini DETR on a mini COCO dataset."""
  
      rng = jax.random.PRNGKey(0)
      np.random.seed(0)
  
      with tfds.testing.mock_data(
          num_examples=50,
          as_dataset_fn=test_util.generate_fake_dataset(
              num_examples=50
          ),
      ):
          _, train_summary, eval_summary = trainer.train_and_evaluate(
              rng=rng,
              config=self.config,
              model_cls=self.model_cls,
              dataset=self.dataset,
              workdir=self.test_dir,
              writer=metric_writers.LoggingWriter(),
          )
  
      expected_summary_keys = [
          "loss_bbox",
          "loss_class",
          "loss_giou",
          "total_loss",
          "object_detection_loss",
      ]
      for i in range(self.config.num_decoder_layers - 1):
          expected_summary_keys += [
              f"loss_class_aux{i}",
              f"loss_bbox_aux{i}",
              f"loss_giou_aux{i}",
          ]
      expected_global_summary_keys = [
          "AP",
          "AP_50",
          "AP_75",
          "AP_small",
          "AP_medium",
          "AP_large",
          "AR_max_1",
          "AR_max_10",
          "AR_max_100",
          "AR_small",
          "AR_medium",
          "AR_large",
      ]
  
      # Check summaries.
      self.assertSameElements(
          expected_summary_keys, train_summary.keys()
      )
      self.assertSameElements(
          expected_summary_keys + expected_global_summary_keys,
          eval_summary.keys(),
      )
  
imports: import shutil
  import tempfile
  from absl.testing import absltest
  from absl.testing import parameterized
  from clu import metric_writers
  from jax import flatten_util
  import jax.numpy
  import jax.random
  import numpy
  from scenic.dataset_lib import datasets
  from scenic.projects.baselines.deformable_detr import trainer
  from scenic.projects.baselines.deformable_detr.configs import mini_config
  from scenic.projects.baselines.deformable_detr.model import DeformableDETRModel
  from scenic.projects.baselines.detr.tests import test_util
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 3 for calling the DETRTrainerTest obj's underlying member method code with example usage: obj.test_deformable_detr_trainer() without return types