Merge pull request #572 from KevinMusgrave/dev

v1.7.1

Author: KevinMusgrave

docs/reducers.md

@@ -122,6 +122,13 @@ reducers.PerAnchorReducer(reducer=None,
    	* ```num_per_row``` is a size N array which indicates how many non-zero losses there are per-row of ```x```.
 
 
+## SumReducer
+This will return the sum of the losses.
+```python
+reducers.SumReducer(**kwargs)
+```
+
+
 ## ThresholdReducer
 This computes the average loss, using only the losses that fall within a specified range.
 

setup.py

@@ -42,7 +42,6 @@
         "scikit-learn",
         "tqdm",
         "torch >= 1.6.0",
-        "torchvision",
     ],
     extras_require={
         "with-hooks": extras_require_with_hooks,

src/pytorch_metric_learning/__init__.py

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

src/pytorch_metric_learning/reducers/__init__.py

@@ -6,4 +6,5 @@
 from .mean_reducer import MeanReducer
 from .multiple_reducers import MultipleReducers
 from .per_anchor_reducer import PerAnchorReducer
+from .sum_reducer import SumReducer
 from .threshold_reducer import ThresholdReducer

src/pytorch_metric_learning/reducers/sum_reducer.py

@@ -0,0 +1,8 @@
+import torch
+
+from pytorch_metric_learning.reducers import MeanReducer
+
+
+class SumReducer(MeanReducer):
+    def element_reduction(self, losses, *_):
+        return torch.sum(losses)

src/pytorch_metric_learning/utils/distributed.py

@@ -30,24 +30,27 @@ def all_gather_embeddings_and_labels(emb, labels):
     if not is_distributed():
         return None, None
     ref_emb = all_gather(emb)
-    ref_labels = all_gather(labels)
+    ref_labels = all_gather(labels) if labels is not None else None
     return ref_emb, ref_labels
 
 
 def gather(emb, labels):
     device = emb.device
-    labels = c_f.to_device(labels, device=device)
+    if labels is not None:
+        labels = c_f.to_device(labels, device=device)
     dist_emb, dist_labels = all_gather_embeddings_and_labels(emb, labels)
     all_emb = torch.cat([emb, dist_emb], dim=0)
-    all_labels = torch.cat([labels, dist_labels], dim=0)
+    all_labels = (
+        torch.cat([labels, dist_labels], dim=0) if dist_labels is not None else None
+    )
     return all_emb, all_labels, labels
 
 
 def gather_emb_and_ref(emb, labels, ref_emb=None, ref_labels=None):
     all_emb, all_labels, labels = gather(emb, labels)
     all_ref_emb, all_ref_labels = None, None
 
-    if ref_emb is not None and ref_labels is not None:
+    if ref_emb is not None:
         all_ref_emb, all_ref_labels, _ = gather(ref_emb, ref_labels)
 
     return all_emb, all_labels, all_ref_emb, all_ref_labels, labels
@@ -81,7 +84,9 @@ def __init__(self, loss, efficient=False):
         self.loss = loss
         self.efficient = efficient
 
-    def forward(self, emb, labels, indices_tuple=None, ref_emb=None, ref_labels=None):
+    def forward(
+        self, emb, labels=None, indices_tuple=None, ref_emb=None, ref_labels=None
+    ):
         world_size = torch.distributed.get_world_size()
         common_args = [emb, labels, indices_tuple, ref_emb, ref_labels, world_size]
         if isinstance(self.loss, CrossBatchMemory):
@@ -99,7 +104,8 @@ def forward_regular_loss(
         )
 
         if self.efficient:
-            all_labels = select_ref_or_regular(all_labels, all_ref_labels)
+            if all_labels is not None:
+                all_labels = select_ref_or_regular(all_labels, all_ref_labels)
             all_emb = select_ref_or_regular(all_emb, all_ref_emb)
             if indices_tuple is None:
                 indices_tuple = get_indices_tuple(labels, all_labels)

tests/reducers/test_mean_reducer.py

@@ -2,14 +2,15 @@
 
 import torch
 
-from pytorch_metric_learning.reducers import MeanReducer
+from pytorch_metric_learning.reducers import MeanReducer, SumReducer
 
 from .. import TEST_DEVICE, TEST_DTYPES
 
 
 class TestMeanReducer(unittest.TestCase):
     def test_mean_reducer(self):
         reducer = MeanReducer()
+        sum_reducer = SumReducer()
         batch_size = 100
         embedding_size = 64
         for dtype in TEST_DTYPES:
@@ -42,3 +43,7 @@ def test_mean_reducer(self):
                 output = reducer(loss_dict, embeddings, labels)
                 correct_output = torch.mean(losses)
                 self.assertTrue(output == correct_output)
+
+                output = sum_reducer(loss_dict, embeddings, labels)
+                correct_output = torch.sum(losses)
+                self.assertTrue(output == correct_output)

tests/utils/test_distributed.py

@@ -64,6 +64,7 @@ def single_process_function(
     miner_fn,
     original_model,
     efficient,
+    pass_labels_to_loss_fn,
 ):
     setup(rank, world_size)
     if TEST_DEVICE == torch.device("cpu"):
@@ -102,9 +103,12 @@ def single_process_function(
         indices_tuple = None
         if miner_fn:
             indices_tuple = miner_fn(outputs, curr_labels, ref_outputs, curr_ref_labels)
-        loss = loss_fn(
-            outputs, curr_labels, indices_tuple, ref_outputs, curr_ref_labels
-        )
+        if miner_fn and not pass_labels_to_loss_fn:
+            loss = loss_fn(outputs, indices_tuple=indices_tuple, ref_emb=ref_outputs)
+        else:
+            loss = loss_fn(
+                outputs, curr_labels, indices_tuple, ref_outputs, curr_ref_labels
+            )
 
         dist.barrier()
         loss.backward()
@@ -162,6 +166,7 @@ def loss_and_miner_tester(
         use_ref,
         loss_kwargs=None,
         miner_kwargs=None,
+        pass_labels_to_loss_fn=True,
     ):
         torch.manual_seed(75210)
         loss_kwargs = {} if loss_kwargs is None else loss_kwargs
@@ -294,6 +299,7 @@ def loss_and_miner_tester(
                         miner_fn,
                         original_model,
                         efficient,
+                        pass_labels_to_loss_fn,
                     ),
                     nprocs=world_size,
                     join=True,
@@ -309,31 +315,35 @@ def test_distributed_tuple_loss(self):
     def test_distributed_tuple_loss_and_miner(self):
         for xbm in [False, True]:
             for use_ref in [False, True]:
-                if xbm and use_ref:
-                    continue
-                self.loss_and_miner_tester(
-                    ContrastiveLoss,
-                    PairMarginMiner,
-                    False,
-                    xbm,
-                    use_ref,
-                    miner_kwargs={"pos_margin": 0.5, "neg_margin": 0.5},
-                )
+                for pass_labels_to_loss_fn in [False, True]:
+                    if xbm and use_ref or xbm and not pass_labels_to_loss_fn:
+                        continue
+                    self.loss_and_miner_tester(
+                        ContrastiveLoss,
+                        PairMarginMiner,
+                        False,
+                        xbm,
+                        use_ref,
+                        miner_kwargs={"pos_margin": 0.5, "neg_margin": 0.5},
+                        pass_labels_to_loss_fn=pass_labels_to_loss_fn,
+                    )
 
     def test_distributed_tuple_loss_efficient(self):
         for use_ref in [False, True]:
             self.loss_and_miner_tester(ContrastiveLoss, None, True, False, use_ref)
 
     def test_distributed_tuple_loss_and_miner_efficient(self):
         for use_ref in [False, True]:
-            self.loss_and_miner_tester(
-                ContrastiveLoss,
-                PairMarginMiner,
-                True,
-                False,
-                use_ref,
-                miner_kwargs={"pos_margin": 0.5, "neg_margin": 0.5},
-            )
+            for pass_labels_to_loss_fn in [False, True]:
+                self.loss_and_miner_tester(
+                    ContrastiveLoss,
+                    PairMarginMiner,
+                    True,
+                    False,
+                    use_ref,
+                    miner_kwargs={"pos_margin": 0.5, "neg_margin": 0.5},
+                    pass_labels_to_loss_fn=pass_labels_to_loss_fn,
+                )
 
     def test_single_proc(self):
         setup(0, 1)