updated tests

tests/extractor/extraction/test_torch_vs_tensorflow.py

@@ -45,21 +45,25 @@ def test_custom_torch_vs_tf_extraction(self):
         pt_model.backend = pt_backend
 
         layer_name = "relu"
-        tf_features = tf_model.extract_features(
-            batches=tf_dl,
-            module_name=layer_name,
-            flatten_acts=False,
-        )
-        pt_features = pt_model.extract_features(
-            batches=pt_dl,
-            module_name=layer_name,
-            flatten_acts=False,
-            output_type="tensor",
-        )
-        expected_features_pt = torch.tensor([[2, 2], [0, 0]])
-        expected_features_tf = np.array([[2, 2], [0, 0]])
-        np.testing.assert_allclose(pt_features, expected_features_pt)
-        np.testing.assert_allclose(tf_features, expected_features_tf)
+        expected_features_pt = torch.tensor([[2., 2.], [0., 0.]])
+        expected_features_tf = np.array([[2., 2.], [0, 0.]])
+
+        for i, batch in enumerate(tf_dl):
+            tf_features = tf_model.extract_batch(
+                batch=batch,
+                module_name=layer_name,
+                flatten_acts=False,
+            )
+            np.testing.assert_allclose(tf_features, expected_features_tf[i][None,:])
+
+        for i, batch in enumerate(pt_dl):
+            pt_features = pt_model.extract_batch(
+                batch=batch,
+                module_name=layer_name,
+                flatten_acts=False,
+                output_type="tensor",
+            )
+            np.testing.assert_allclose(pt_features, expected_features_pt[i][None,:])
 
         layer_name = "relu2"
         expected_features = np.array([[4., 4.], [0., 0.]])

tests/helper.py

@@ -100,6 +100,12 @@
         "pretrained": True,
         "source": "keras",
     },
+    "VGG19_keras": {
+        "model_name": "VGG19",
+        "modules": ["block1_conv1", "flatten"],
+        "pretrained": False,
+        "source": "keras",
+    },
     # Vissl models
     "simclr-rn50": {
         "model_name": "simclr-rn50",

thingsvision/core/extraction/base.py

@@ -105,8 +105,7 @@ def extract_batch(
             Returns the feature matrix (e.g., $X \in \mathbb{R}^{B \times d}$ if penultimate or logits layer or flatten_acts = True).
         """
         raise NotImplementedError
-
-
+
     @abc.abstractmethod
     def _extract_batch(
         self,

thingsvision/core/extraction/tensorflow.py

@@ -34,7 +34,7 @@ def __init__(
         if not self.model:
             self.load_model()
         self.prepare_inference()
-        
+
     def _extract_batch(
         self,
         batch: Array,

thingsvision/core/extraction/torch.py

@@ -86,7 +86,7 @@ def extract_batch(
             act = self._to_numpy(act)
         self._unregister_hook()
         return act
-        
+
     @torch.no_grad()
     def _extract_batch(
         self,