SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[DETR, YOLOS] Fix device bug (#21974)

Browse Source 
* Fix integration test

* Add test

* Add test
This commit is contained in:
NielsRogge
2023-03-07 13:34:04 +01:00
committed by GitHub
parent eec46b4f75
commit 95408e9953
4 changed files with 52 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
tests/models/yolos/test_modeling_yolos.py
View File

@@ -360,7 +360,7 @@ class YolosModelIntegrationTest(unittest.TestCase):
with torch.no_grad():
outputs = model(inputs.pixel_values)
# verify the logits
# verify outputs
expected_shape = torch.Size((1, 100, 92))
self.assertEqual(outputs.logits.shape, expected_shape)
@@ -373,3 +373,16 @@ class YolosModelIntegrationTest(unittest.TestCase):
)
self.assertTrue(torch.allclose(outputs.logits[0, :3, :3], expected_slice_logits, atol=1e-4))
self.assertTrue(torch.allclose(outputs.pred_boxes[0, :3, :3], expected_slice_boxes, atol=1e-4))
# verify postprocessing
results = feature_extractor.post_process_object_detection(
outputs, threshold=0.3, target_sizes=[image.size[::-1]]
)[0]
expected_scores = torch.tensor([0.9994, 0.9790, 0.9964, 0.9972, 0.9861]).to(torch_device)
expected_labels = [75, 75, 17, 63, 17]
expected_slice_boxes = torch.tensor([335.0609, 79.3848, 375.4216, 187.2495]).to(torch_device)
self.assertEqual(len(results["scores"]), 5)
self.assertTrue(torch.allclose(results["scores"], expected_scores, atol=1e-4))
self.assertSequenceEqual(results["labels"].tolist(), expected_labels)
self.assertTrue(torch.allclose(results["boxes"][0, :], expected_slice_boxes))