fix auxiliary loss training in DetrSegmentation (#28354)

* fix auxiliary loss training in detrSegmentation

* add auxiliary_loss testing
This commit is contained in:
Sangbum Daniel Choi
2024-01-09 19:17:07 +09:00
committed by GitHub
parent 8604dd308d
commit 357971ec36
2 changed files with 19 additions and 3 deletions

View File

@@ -399,6 +399,22 @@ class DetrModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin
self.assertIsNotNone(decoder_attentions.grad)
self.assertIsNotNone(cross_attentions.grad)
def test_forward_auxiliary_loss(self):
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
config.auxiliary_loss = True
# only test for object detection and segmentation model
for model_class in self.all_model_classes[1:]:
model = model_class(config)
model.to(torch_device)
inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True)
outputs = model(**inputs)
self.assertIsNotNone(outputs.auxiliary_outputs)
self.assertEqual(len(outputs.auxiliary_outputs), self.model_tester.num_hidden_layers - 1)
def test_forward_signature(self):
config, _ = self.model_tester.prepare_config_and_inputs_for_common()