device agnostic models testing (#27146)

* device agnostic models testing

* add decorator `require_torch_fp16`

* make style

* apply review suggestion

* Oops, the fp16 decorator was misused
This commit is contained in:
Hz, Ji
2023-11-01 01:12:14 +08:00
committed by GitHub
parent 77930f8a01
commit 50378cbf6c
51 changed files with 369 additions and 154 deletions

View File

@@ -25,7 +25,7 @@ from transformers.file_utils import cached_property
from transformers.testing_utils import (
require_timm,
require_torch,
require_torch_gpu,
require_torch_accelerator,
require_vision,
slow,
torch_device,
@@ -648,7 +648,7 @@ class DeformableDetrModelIntegrationTests(unittest.TestCase):
self.assertEqual(outputs.pred_boxes.shape, expected_shape_boxes)
self.assertTrue(torch.allclose(outputs.pred_boxes[0, :3, :3], expected_boxes, atol=1e-4))
@require_torch_gpu
@require_torch_accelerator
def test_inference_object_detection_head_equivalence_cpu_gpu(self):
image_processor = self.default_image_processor
image = prepare_img()
@@ -663,10 +663,10 @@ class DeformableDetrModelIntegrationTests(unittest.TestCase):
cpu_outputs = model(pixel_values, pixel_mask)
# 2. run model on GPU
model.to("cuda")
model.to(torch_device)
with torch.no_grad():
gpu_outputs = model(pixel_values.to("cuda"), pixel_mask.to("cuda"))
gpu_outputs = model(pixel_values.to(torch_device), pixel_mask.to(torch_device))
# 3. assert equivalence
for key in cpu_outputs.keys():