class_docstr: Unit tests for detr test_input_pipeline_detection.py.
docstr: Tests dataset builder for coco_deformable_detr_detection.
code: def test_dataset_builder_coco_deformable_detr_detection(self):
      """Tests dataset builder for coco_deformable_detr_detection."""
      num_shards = jax.local_device_count()
      batch_size = num_shards * 2
      eval_batch_size = num_shards * 2
  
      dataset_config = ml_collections.ConfigDict()
      dataset_config.max_size = 1333
      dataset_config.valid_max_size = 1333
  
      with tfds.testing.mock_data(
          num_examples=50,
          as_dataset_fn=test_util.generate_fake_dataset(
              num_examples=50
          ),
      ):
          dataset = input_pipeline_detection.get_dataset(
              batch_size=batch_size,
              eval_batch_size=eval_batch_size,
              dataset_configs=dataset_config,
              num_shards=num_shards,
          )
  
          # A dataset should at least provide `train_iter` and `valid_iter`.
          self.assertIsNotNone(dataset.train_iter)
          self.assertIsNotNone(dataset.valid_iter)
  
          train_batch = next(dataset.train_iter)
          eval_batch = next(dataset.valid_iter)
  
      # Check shapes.
      # Tests first two shape dimensions.
      expected_shape = [
          num_shards,
          batch_size // num_shards,
          1333,
          1333,
          3,
      ]
      expected_shape_eval = [
          num_shards,
          eval_batch_size // num_shards,
          1333,
          1333,
          3,
      ]
      self.assertSequenceEqual(
          train_batch["inputs"].shape, expected_shape
      )
      self.assertSequenceEqual(
          eval_batch["inputs"].shape, expected_shape_eval
      )
  
      self.assertEqual(
          train_batch["inputs"].shape[:-1],
          train_batch["padding_mask"].shape,
      )
      self.assertEqual(
          eval_batch["inputs"].shape[:-1],
          eval_batch["padding_mask"].shape,
      )
  
file_path: scenic/scenic/projects/baselines/deformable_detr/tests/test_input_pipeline.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import ml_collections
  from scenic.projects.baselines.deformable_detr import input_pipeline_detection
  from scenic.projects.baselines.detr.tests import test_util
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling the DeformableDETRInputPipelineTests obj's underlying member method code with example usage: obj.test_dataset_builder_coco_deformable_detr_detection() without return types