Add test_export for dpt and zoedepth

Signed-off-by: Phillip Kuznetsov <[email protected]>

Author: philkuz

tests/models/dpt/test_modeling_dpt.py

@@ -18,6 +18,7 @@
 
 from transformers import DPTConfig
 from transformers.file_utils import is_torch_available, is_vision_available
+from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_4
 from transformers.testing_utils import require_torch, require_vision, slow, torch_device
 
 from ...test_configuration_common import ConfigTester
@@ -410,3 +411,24 @@ def test_post_processing_depth_estimation(self):
         ).squeeze()
         self.assertTrue(output_enlarged.shape == expected_shape)
         self.assertTrue(torch.allclose(predicted_depth_l, output_enlarged, rtol=1e-3))
+
+    def test_export(self):
+        for strict in [True, False]:
+            with self.subTest(strict=strict):
+                if not is_torch_greater_or_equal_than_2_4:
+                    self.skipTest(reason="This test requires torch >= 2.4 to run.")
+                model = DPTForSemanticSegmentation.from_pretrained("Intel/dpt-large-ade").to(torch_device).eval()
+                image_processor = DPTImageProcessor.from_pretrained("Intel/dpt-large-ade")
+                image = prepare_img()
+                inputs = image_processor(images=image, return_tensors="pt").to(torch_device)
+
+                exported_program = torch.export.export(
+                    model,
+                    args=(inputs["pixel_values"],),
+                    strict=strict,
+                )
+                with torch.no_grad():
+                    eager_outputs = model(**inputs)
+                    exported_outputs = exported_program.module().forward(inputs["pixel_values"])
+                self.assertEqual(eager_outputs.logits.shape, exported_outputs.logits.shape)
+                self.assertTrue(torch.allclose(eager_outputs.logits, exported_outputs.logits, atol=1e-4))

tests/models/zoedepth/test_modeling_zoedepth.py

@@ -20,6 +20,7 @@
 
 from transformers import Dinov2Config, ZoeDepthConfig
 from transformers.file_utils import is_torch_available, is_vision_available
+from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_4
 from transformers.testing_utils import require_torch, require_vision, slow, torch_device
 
 from ...test_configuration_common import ConfigTester
@@ -354,3 +355,29 @@ def test_inference_depth_estimation_post_processing_pad_flip(self):
         model = ZoeDepthForDepthEstimation.from_pretrained("Intel/zoedepth-nyu-kitti").to(torch_device)
 
         self.check_post_processing_test(image_processor, images, model, pad_input=True, flip_aug=True)
+
+    def test_export(self):
+        self.skipTest(
+            reason="This test fails because the beit backbone of ZoeDepth is not compatible with torch.export"
+        )
+        for strict in [True, False]:
+            with self.subTest(strict=True):
+                if not is_torch_greater_or_equal_than_2_4:
+                    self.skipTest(reason="This test requires torch >= 2.4 to run.")
+                model = ZoeDepthForDepthEstimation.from_pretrained("Intel/zoedepth-nyu").to(torch_device).eval()
+                image_processor = ZoeDepthImageProcessor.from_pretrained("Intel/zoedepth-nyu")
+                image = prepare_img()
+                inputs = image_processor(images=image, return_tensors="pt").to(torch_device)
+
+                exported_program = torch.export.export(
+                    model,
+                    args=(inputs["pixel_values"],),
+                    strict=strict,
+                )
+                with torch.no_grad():
+                    eager_outputs = model(**inputs)
+                    exported_outputs = exported_program.module().forward(inputs["pixel_values"])
+                self.assertEqual(eager_outputs.predicted_depth.shape, exported_outputs.predicted_depth.shape)
+                self.assertTrue(
+                    torch.allclose(eager_outputs.predicted_depth, exported_outputs.predicted_depth, atol=1e-4)
+                )