Added VGG presets test cases (#2090)

keras-team · Feb 11, 2025 · 0f7afd7 · 0f7afd7
1 parent be43414
commit 0f7afd7
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 0 deletions.
diff --git a/keras_hub/src/models/vgg/vgg_backbone_test.py b/keras_hub/src/models/vgg/vgg_backbone_test.py
@@ -30,3 +30,15 @@ def test_saved_model(self):
             init_kwargs=self.init_kwargs,
             input_data=self.input_data,
         )
+
+    @pytest.mark.extra_large
+    def test_all_presets(self):
+        # we need at least 32x32 image resolution here to satisfy the presets'
+        # downsampling constraints
+        images = np.ones((2, 32, 32, 3), dtype="float32")
+        for preset in VGGBackbone.presets:
+            self.run_preset_test(
+                cls=VGGBackbone,
+                preset=preset,
+                input_data=images,
+            )
diff --git a/keras_hub/src/models/vgg/vgg_image_classifier_test.py b/keras_hub/src/models/vgg/vgg_image_classifier_test.py
@@ -51,3 +51,17 @@ def test_saved_model(self):
             init_kwargs=self.init_kwargs,
             input_data=self.images,
         )
+
+    @pytest.mark.extra_large
+    def test_all_presets(self):
+        # we need at least 32x32 image resolution here to satisfy the presets'
+        # downsampling constraints
+        images = np.ones((2, 32, 32, 3), dtype="float32")
+        for preset in VGGImageClassifier.presets:
+            self.run_preset_test(
+                cls=VGGImageClassifier,
+                preset=preset,
+                init_kwargs={"num_classes": 2},
+                input_data=images,
+                expected_output_shape=(2, 2),
+            )