def build_model(self):
"""Build DETR model."""
model = detr.DETR(self._task_config.num_queries,
self._task_config.num_hidden,
self._task_config.num_classes,
self._task_config.num_encoder_layers,
self._task_config.num_decoder_layers)
return model
def test_forward(self):
num_queries = 10
hidden_size = 128
num_classes = 10
image_size = 640
batch_size = 2
model = detr.DETR(num_queries, hidden_size, num_classes)
outs = model(tf.ones((batch_size, image_size, image_size, 3)))
self.assertLen(outs, 6) # intermediate decoded outputs.
for out in outs:
self.assertAllEqual(tf.shape(out['cls_outputs']),
(batch_size, num_queries, num_classes))
self.assertAllEqual(tf.shape(out['box_outputs']),
(batch_size, num_queries, 4))
def build_model(self):
"""Build DETR model."""
input_specs = tf.keras.layers.InputSpec(
shape=[None] + self._task_config.model.input_size)
backbone = backbones.factory.build_backbone(
input_specs=input_specs,
backbone_config=self._task_config.model.backbone,
norm_activation_config=self._task_config.model.norm_activation)
model = detr.DETR(backbone, self._task_config.model.num_queries,
self._task_config.model.hidden_size,
self._task_config.model.num_classes,
self._task_config.model.num_encoder_layers,
self._task_config.model.num_decoder_layers)
return model
def test_forward(self):
num_queries = 10
hidden_size = 128
num_classes = 10
image_size = 640
batch_size = 2
backbone = resnet.ResNet(50, bn_trainable=False)
backbone_endpoint_name = '5'
model = detr.DETR(backbone, backbone_endpoint_name, num_queries,
hidden_size, num_classes)
outs = model(tf.ones((batch_size, image_size, image_size, 3)))
self.assertLen(outs, 6) # intermediate decoded outputs.
for out in outs:
self.assertAllEqual(tf.shape(out['cls_outputs']),
(batch_size, num_queries, num_classes))
self.assertAllEqual(tf.shape(out['box_outputs']),
(batch_size, num_queries, 4))