Merge pull request #574 from KevinMusgrave/dev

v1.7.2

Author: KevinMusgrave

src/pytorch_metric_learning/__init__.py

@@ -1 +1 @@
-__version__ = "1.7.1"
+__version__ = "1.7.2"

src/pytorch_metric_learning/distances/base_distance.py

@@ -16,6 +16,7 @@ def __init__(
 
     def forward(self, query_emb, ref_emb=None):
         self.reset_stats()
+        self.check_shapes(query_emb, ref_emb)
         query_emb_normalized = self.maybe_normalize(query_emb)
         if ref_emb is None:
             ref_emb = query_emb
@@ -87,3 +88,9 @@ def set_stats(self, stats_dict):
         for k, v in stats_dict.items():
             self.add_to_recordable_attributes(name=k, is_stat=True)
             setattr(self, k, v)
+
+    def check_shapes(self, query_emb, ref_emb):
+        if query_emb.ndim != 2 or (ref_emb is not None and ref_emb.ndim != 2):
+            raise ValueError(
+                "embeddings must be a 2D tensor of shape (batch_size, embedding_size)"
+            )

src/pytorch_metric_learning/utils/common_functions.py

@@ -412,10 +412,6 @@ def return_input(x):
 def check_shapes(embeddings, labels):
     if labels is not None and embeddings.shape[0] != labels.shape[0]:
         raise ValueError("Number of embeddings must equal number of labels")
-    if embeddings.ndim != 2:
-        raise ValueError(
-            "embeddings must be a 2D tensor of shape (batch_size, embedding_size)"
-        )
     if labels is not None and labels.ndim != 1:
         raise ValueError("labels must be a 1D tensor of shape (batch_size,)")
 

tests/distances/test_custom_check_shape.py

@@ -0,0 +1,27 @@
+import unittest
+
+import torch
+
+from pytorch_metric_learning.distances import BaseDistance, LpDistance
+from pytorch_metric_learning.losses import TripletMarginLoss
+
+
+class CustomDistance(BaseDistance):
+    def compute_mat(self, query_emb, ref_emb):
+        return torch.randn(query_emb.shape[0], ref_emb.shape[0])
+
+    def check_shapes(self, query_emb, ref_emb):
+        pass
+
+
+class TestCustomCheckShape(unittest.TestCase):
+    def test_custom_embedding_ndim(self):
+        embeddings = torch.randn(32, 3, 128)
+        labels = torch.randint(0, 10, size=(32,))
+        loss_fn = TripletMarginLoss(distance=LpDistance())
+
+        with self.assertRaises(ValueError):
+            loss_fn(embeddings, labels)
+
+        loss_fn = TripletMarginLoss(distance=CustomDistance())
+        loss_fn(embeddings, labels)

tests/utils/test_common_functions.py

@@ -69,7 +69,7 @@ def test_collect_stats_flag(self):
 
     def test_check_shapes(self):
         embeddings = torch.randn(32, 512, 3)
-        labels = torch.randn(32)
+        labels = torch.randint(0, 10, size=(32,))
         loss_fn = TripletMarginLoss()
 
         # embeddings is 3-dimensional