Fix RT-DETR inference with float16 and bfloat16 (#31639)

* [run_slow] rt_detr

* Fix positional embeddings and anchors dtypes

* [run slow] rt_detr

* Apply suggestions from code review

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

* Fixup

---------

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>
This commit is contained in:
Pavel Iakubovskii
2024-06-26 17:50:10 +01:00
committed by GitHub
parent 3f93fd0694
commit b1ec745475
2 changed files with 29 additions and 4 deletions

View File

@@ -18,6 +18,8 @@ import inspect
import math
import unittest
from parameterized import parameterized
from transformers import (
RTDetrConfig,
RTDetrImageProcessor,
@@ -25,7 +27,7 @@ from transformers import (
is_torch_available,
is_vision_available,
)
from transformers.testing_utils import require_torch, require_vision, torch_device
from transformers.testing_utils import require_torch, require_torch_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property
from ...test_configuration_common import ConfigTester
@@ -606,6 +608,28 @@ class RTDetrModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
msg=f"Parameter {name} of model {model_class} seems not properly initialized",
)
@parameterized.expand(["float32", "float16", "bfloat16"])
@require_torch_gpu
@slow
def test_inference_with_different_dtypes(self, torch_dtype_str):
torch_dtype = {
"float32": torch.float32,
"float16": torch.float16,
"bfloat16": torch.bfloat16,
}[torch_dtype_str]
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
for model_class in self.all_model_classes:
model = model_class(config)
model.to(torch_device).to(torch_dtype)
model.eval()
for key, tensor in inputs_dict.items():
if tensor.dtype == torch.float32:
inputs_dict[key] = tensor.to(torch_dtype)
with torch.no_grad():
_ = model(**self._prepare_for_class(inputs_dict, model_class))
TOLERANCE = 1e-4