Fix DINOV2 test.

Author: james77777778

keras_hub/src/models/dinov2/dinov2_backbone.py

@@ -19,6 +19,10 @@ class DINOV2Backbone(FeaturePyramidBackbone):
     DINOV2 model with any number of layers, heads, and embedding dimensions. To
     load preset architectures and weights, use the `from_preset` constructor.
 
+    Note that this backbone is a Feature Pyramid Backbone that can output
+    intermediate feature maps from different stages of the model. See the
+    example below for how to access these feature pyramid outputs.
+
     Note that this backbone supports interpolation of the position embeddings
     to the input image shape. This is useful when the input image shape is
     different from the shape used to train the position embeddings. The
@@ -97,6 +101,16 @@ class DINOV2Backbone(FeaturePyramidBackbone):
         position_embedding_shape=(518, 518),
     )
     model(input_data)
+
+    # Accessing feature pyramid outputs.
+    backbone = keras_hub.models.DINOV2Backbone.from_preset(
+        "dinov2_base", image_shape=(224, 224, 3)
+    )
+    model = keras.Model(
+        inputs=backbone.inputs,
+        outputs=backbone.pyramid_outputs,
+    )
+    features = model(input_data)
     ```
     """
 

keras_hub/src/models/dinov2/dinov2_backbone_test.py

@@ -19,22 +19,26 @@ def setUp(self):
             "layer_scale_init_value": 1.0,
             "num_register_tokens": 0,
             "use_swiglu_ffn": False,
-            "image_shape": (64, 64, 3),
+            "image_shape": (70, 70, 3),
+            "name": "dinov2_backbone",
         }
         self.input_data = {
-            "images": ops.ones((2, 64, 64, 3)),
+            "images": ops.ones((2, 70, 70, 3)),
         }
 
     def test_backbone_basics(self):
         patch_size = self.init_kwargs["patch_size"]
         image_size = self.init_kwargs["image_shape"][0]
         hidden_dim = self.init_kwargs["hidden_dim"]
         sequence_length = (image_size // patch_size) ** 2 + 1
-        self.run_backbone_test(
+        self.run_vision_backbone_test(
             cls=DINOV2Backbone,
             init_kwargs=self.init_kwargs,
             input_data=self.input_data,
             expected_output_shape=(2, sequence_length, hidden_dim),
+            expected_pyramid_output_keys=["Stem", "Stage1", "Stage2"],
+            expected_pyramid_image_sizes=[(sequence_length, hidden_dim)] * 3,
+            run_data_format_check=False,
         )
 
     @pytest.mark.large
@@ -107,10 +111,11 @@ def setUp(self):
             "layer_scale_init_value": 1.0,
             "num_register_tokens": 4,
             "use_swiglu_ffn": True,
-            "image_shape": (64, 64, 3),
+            "image_shape": (70, 70, 3),
+            "name": "dinov2_backbone",
         }
         self.input_data = {
-            "images": ops.ones((2, 64, 64, 3)),
+            "images": ops.ones((2, 70, 70, 3)),
         }
 
     def test_backbone_basics(self):
@@ -121,11 +126,14 @@ def test_backbone_basics(self):
         sequence_length = (
             (image_size // patch_size) ** 2 + 1 + num_register_tokens
         )
-        self.run_backbone_test(
+        self.run_vision_backbone_test(
             cls=DINOV2Backbone,
             init_kwargs=self.init_kwargs,
             input_data=self.input_data,
             expected_output_shape=(2, sequence_length, hidden_dim),
+            expected_pyramid_output_keys=["Stem", "Stage1", "Stage2"],
+            expected_pyramid_image_sizes=[(sequence_length, hidden_dim)] * 3,
+            run_data_format_check=False,
         )
 
     @pytest.mark.large

keras_hub/src/tests/test_case.py

@@ -538,10 +538,11 @@ def run_vision_backbone_test(
 
             self.assertIsInstance(output_data, dict)
             self.assertEqual(
-                list(output_data.keys()), list(backbone.pyramid_outputs.keys())
+                sorted(output_data.keys()),
+                sorted(backbone.pyramid_outputs.keys()),
             )
             self.assertEqual(
-                list(output_data.keys()), expected_pyramid_output_keys
+                sorted(output_data.keys()), sorted(expected_pyramid_output_keys)
             )
             # check height and width of each level.
             for i, (k, v) in enumerate(output_data.items()):