Merge pull request #577 from KevinMusgrave/dev

v1.7.3

Author: KevinMusgrave

src/pytorch_metric_learning/__init__.py

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

src/pytorch_metric_learning/losses/cross_batch_memory.py

@@ -18,6 +18,8 @@ def __init__(self, loss, embedding_size, memory_size=1024, miner=None, **kwargs)
         )
 
     def forward(self, embeddings, labels, indices_tuple=None, enqueue_idx=None):
+        if indices_tuple is not None and enqueue_idx is not None:
+            raise ValueError("indices_tuple and enqueue_idx are mutually exclusive")
         if enqueue_idx is not None:
             assert len(enqueue_idx) <= len(self.embedding_memory)
             assert len(enqueue_idx) < len(embeddings)
@@ -46,7 +48,6 @@ def forward(self, embeddings, labels, indices_tuple=None, enqueue_idx=None):
             labels_for_queue = labels
             do_remove_self_comparisons = True
 
-        batch_size = len(embeddings)
         queue_batch_size = len(emb_for_queue)
         self.add_to_memory(emb_for_queue, labels_for_queue, queue_batch_size)
 
@@ -58,7 +59,6 @@ def forward(self, embeddings, labels, indices_tuple=None, enqueue_idx=None):
             L_mem = self.label_memory
 
         indices_tuple = self.create_indices_tuple(
-            batch_size,
             embeddings,
             labels,
             E_mem,
@@ -85,7 +85,6 @@ def add_to_memory(self, embeddings, labels, batch_size):
 
     def create_indices_tuple(
         self,
-        batch_size,
         embeddings,
         labels,
         E_mem,
@@ -117,7 +116,9 @@ def create_indices_tuple(
         return indices_tuple
 
     def reset_queue(self):
-        self.embedding_memory = torch.zeros(self.memory_size, self.embedding_size)
-        self.label_memory = torch.zeros(self.memory_size).long()
+        self.register_buffer(
+            "embedding_memory", torch.zeros(self.memory_size, self.embedding_size)
+        )
+        self.register_buffer("label_memory", torch.zeros(self.memory_size).long())
         self.has_been_filled = False
         self.queue_idx = 0

src/pytorch_metric_learning/utils/distributed.py

@@ -66,6 +66,19 @@ def get_indices_tuple(labels, ref_labels, embeddings=None, ref_emb=None, miner=N
     return lmu.remove_self_comparisons(indices_tuple, curr_batch_idx, len(ref_labels))
 
 
+def get_corrected_enqueue_idx(enqueue_idx, emb):
+    if enqueue_idx is None:
+        return enqueue_idx
+    enqueue_idx = c_f.to_device(enqueue_idx, device=emb.device)
+    bs = len(emb)
+    e_len = len(enqueue_idx)
+    world_size = torch.distributed.get_world_size()
+    enqueue_idx = torch.cat([enqueue_idx, all_gather(enqueue_idx)], dim=0)
+    for i in range(e_len, e_len * world_size, e_len):
+        enqueue_idx[i:] += bs
+    return enqueue_idx
+
+
 def select_ref_or_regular(regular, ref):
     return regular if ref is None else ref
 
@@ -85,12 +98,18 @@ def __init__(self, loss, efficient=False):
         self.efficient = efficient
 
     def forward(
-        self, emb, labels=None, indices_tuple=None, ref_emb=None, ref_labels=None
+        self,
+        emb,
+        labels=None,
+        indices_tuple=None,
+        ref_emb=None,
+        ref_labels=None,
+        enqueue_idx=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):
-            return self.forward_cross_batch(*common_args)
+            return self.forward_cross_batch(*common_args, enqueue_idx)
         return self.forward_regular_loss(*common_args)
 
     def forward_regular_loss(
@@ -118,20 +137,28 @@ def forward_regular_loss(
         return loss * world_size
 
     def forward_cross_batch(
-        self, emb, labels, indices_tuple, ref_emb, ref_labels, world_size
+        self,
+        emb,
+        labels,
+        indices_tuple,
+        ref_emb,
+        ref_labels,
+        world_size,
+        enqueue_idx,
     ):
         if ref_emb is not None or ref_labels is not None:
             raise ValueError(
                 "CrossBatchMemory is not compatible with ref_emb and ref_labels"
             )
 
         if world_size <= 1:
-            return self.loss(emb, labels, indices_tuple)
+            return self.loss(emb, labels, indices_tuple, enqueue_idx)
 
         all_emb, all_labels, _, _, _ = gather_emb_and_ref(
             emb, labels, ref_emb, ref_labels
         )
-        loss = self.loss(all_emb, all_labels, indices_tuple)
+        enqueue_idx = get_corrected_enqueue_idx(enqueue_idx, emb)
+        loss = self.loss(all_emb, all_labels, indices_tuple, enqueue_idx)
         return loss * world_size
 
 

tests/losses/test_cross_batch_memory.py

@@ -501,7 +501,6 @@ def test_input_indices_tuple(self):
                         indices_tuple, self.loss.curr_batch_idx, self.loss.memory_size
                     )
                     a1, p, a2, n = self.loss.create_indices_tuple(
-                        batch_size,
                         embeddings,
                         labels,
                         self.loss.embedding_memory,
@@ -520,6 +519,26 @@ def test_input_indices_tuple(self):
                     self.assertTrue(torch.all(a2 == torch.cat([a2i, a2ii])))
                     self.assertTrue(torch.all(n == torch.cat([ni, nii])))
 
+    def test_reset_queue(self):
+        self.loss = CrossBatchMemory(
+            loss=ContrastiveLoss(),
+            embedding_size=self.embedding_size,
+            memory_size=self.memory_size,
+        )
+
+        init_emb = torch.zeros(self.memory_size, self.embedding_size)
+        init_label = torch.zeros(self.memory_size).long()
+        self.assertTrue(torch.equal(self.loss.embedding_memory, init_emb))
+        self.assertTrue(torch.equal(self.loss.label_memory, init_label))
+
+        self.loss(torch.randn(32, 128), torch.randint(0, 2, size=(32,)))
+        self.assertTrue(not torch.equal(self.loss.embedding_memory, init_emb))
+        self.assertTrue(not torch.equal(self.loss.label_memory, init_label))
+
+        self.loss.reset_queue()
+        self.assertTrue(torch.equal(self.loss.embedding_memory, init_emb))
+        self.assertTrue(torch.equal(self.loss.label_memory, init_label))
+
 
 if __name__ == "__main__":
     unittest.main()

tests/utils/test_distributed.py

@@ -65,6 +65,8 @@ def single_process_function(
     original_model,
     efficient,
     pass_labels_to_loss_fn,
+    use_xbm_enqueue_idx,
+    enqueue_idx,
 ):
     setup(rank, world_size)
     if TEST_DEVICE == torch.device("cpu"):
@@ -105,6 +107,10 @@ def single_process_function(
             indices_tuple = miner_fn(outputs, curr_labels, 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)
+        elif use_xbm_enqueue_idx and isinstance(loss_fn.loss, CrossBatchMemory):
+            loss = loss_fn(
+                outputs, curr_labels, indices_tuple, enqueue_idx=enqueue_idx[rank]
+            )
         else:
             loss = loss_fn(
                 outputs, curr_labels, indices_tuple, ref_outputs, curr_ref_labels
@@ -149,6 +155,22 @@ def create_labels(batch_size, world_size, iterations):
     ]
 
 
+def create_enqueue_idx(batch_size, world_size):
+    # enqueue every other embedding
+    local_enqueue_idx = [
+        (torch.randint(0, batch_size, size=(batch_size // 4,))).long()
+        for _ in range(world_size)
+    ]
+    global_enqueue_idx = []
+    for i, x in enumerate(local_enqueue_idx):
+        if i == 0:
+            global_enqueue_idx.append(x)
+        else:
+            global_enqueue_idx.append(x + batch_size)
+    global_enqueue_idx = torch.cat(global_enqueue_idx, dim=0)
+    return local_enqueue_idx, global_enqueue_idx
+
+
 def get_all_outputs_and_labels(inputs, labels, model, iteration):
     all_inputs = torch.cat(inputs[iteration], dim=0).to(TEST_DEVICE)
     all_labels = torch.cat(labels[iteration], dim=0).to(TEST_DEVICE)
@@ -167,6 +189,7 @@ def loss_and_miner_tester(
         loss_kwargs=None,
         miner_kwargs=None,
         pass_labels_to_loss_fn=True,
+        use_xbm_enqueue_idx=False,
     ):
         torch.manual_seed(75210)
         loss_kwargs = {} if loss_kwargs is None else loss_kwargs
@@ -222,6 +245,10 @@ def loss_and_miner_tester(
                     )
                     ref_labels = create_labels(batch_size, world_size, iterations)
 
+                local_enqueue_idx, global_enqueue_idx = create_enqueue_idx(
+                    batch_size, world_size
+                )
+
                 for aaa in range(iterations):
                     optimizer.zero_grad()
                     all_outputs, all_labels = get_all_outputs_and_labels(
@@ -269,8 +296,11 @@ def loss_and_miner_tester(
                                 all_outputs, all_labels, all_ref_outputs, all_ref_labels
                             )
                         if xbm:
+                            enqueue_idx = (
+                                global_enqueue_idx if use_xbm_enqueue_idx else None
+                            )
                             loss = original_loss_fn(
-                                all_outputs, all_labels, indices_tuple
+                                all_outputs, all_labels, indices_tuple, enqueue_idx
                             )
                         else:
                             loss = original_loss_fn(
@@ -300,6 +330,8 @@ def loss_and_miner_tester(
                         original_model,
                         efficient,
                         pass_labels_to_loss_fn,
+                        use_xbm_enqueue_idx,
+                        local_enqueue_idx,
                     ),
                     nprocs=world_size,
                     join=True,
@@ -308,9 +340,17 @@ def loss_and_miner_tester(
     def test_distributed_tuple_loss(self):
         for xbm in [False, True]:
             for use_ref in [False, True]:
-                if xbm and use_ref:
-                    continue
-                self.loss_and_miner_tester(ContrastiveLoss, None, False, xbm, use_ref)
+                for use_xbm_enqueue_idx in [False, True]:
+                    if xbm and use_ref:
+                        continue
+                    self.loss_and_miner_tester(
+                        ContrastiveLoss,
+                        None,
+                        False,
+                        xbm,
+                        use_ref,
+                        use_xbm_enqueue_idx=use_xbm_enqueue_idx,
+                    )
 
     def test_distributed_tuple_loss_and_miner(self):
         for xbm in [False, True]: