class_docstr: Test DETRModel.
class_docstr_tok: ["Test","DETRModel","."]
code: def setUp(self):
      super(DETRModelTest, self).setUp()
  
      self.num_classes = 5
      self.input_shape = (3, 128, 128, 3)
  
      # create and initialize the model
      model_cls = model.DETRModel
      config = ml_collections.ConfigDict(
          dict(
              hidden_dim=32,
              num_queries=8,
              query_emb_size=None,
              transformer_num_heads=2,
              transformer_num_encoder_layers=1,
              transformer_qkv_dim=32,
              transformer_mlp_dim=32,
              transformer_normalize_before=False,
              backbone_num_filters=32,
              backbone_num_layers=14,
              dropout_rate=0.1,
              attention_dropout_rate=0.1,
              class_loss_coef=1.0,
              bbox_loss_coef=1.0,
              giou_loss_coef=1.0,
              eos_coef=1.0,
              aux_loss=False,
              transformer_num_decoder_layers=3,
          )
      )
      self.model = model_cls(
          config=config,
          dataset_meta_data={
              "num_classes": self.num_classes,
              "target_is_onehot": False,
          },
      )
      rng = random.PRNGKey(0)
      initial_params = self.model.flax_model.init(
          rng, jnp.zeros(self.input_shape, jnp.float32), train=False
      )
      flax_model = functools.partial(
          self.model.flax_model.apply,
          initial_params,
          mutable=["intermediates"],
          capture_intermediates=lambda mdl, _: mdl.name
          == "attn_weights",
      )
  
      # a fake batch with 3 examples
      self.batch = {
          "inputs": jnp.array(
              np.random.normal(size=self.input_shape)
          ).astype(jnp.float32),
          "padding_mask": jnp.array(
              np.random.normal(size=self.input_shape[:-1])
          ).astype(jnp.float32),
          "label": {
              "labels": jnp.array(
                  np.random.randint(
                      self.num_classes,
                      size=(3, self.model.config.num_queries),
                  )
              ),
              "boxes": jnp.array(
                  np.random.uniform(
                      size=(3, self.model.config.num_queries, 4),
                      low=0.0,
                      high=1.0,
                  ),
                  dtype=jnp.float32,
              ),
          },
      }
  
      self.outputs, self.variables = flax_model(
          self.batch["inputs"],
          padding_mask=self.batch["padding_mask"],
          train=False,
      )
  
      seq = np.arange(self.model.config.num_queries, dtype=np.int32)
      seq_rev = seq[::-1]
      seq_21 = np.concatenate(
          [
              seq[self.model.config.num_queries // 2 :],
              seq[: self.model.config.num_queries // 2],
          ]
      )
      self.indices = jnp.array(
          [(seq, seq_rev), (seq_rev, seq), (seq, seq_21)]
      )
  
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 1 for calling the DETRModelTest obj's underlying member method code with example usage: obj.setUp() without return types