Merge pull request #449 from KevinMusgrave/dev

v1.2.1

Author: Kevin Musgrave

src/pytorch_metric_learning/__init__.py

@@ -1 +1 @@
-__version__ = "1.2.0"
+__version__ = "1.2.1"

src/pytorch_metric_learning/utils/distributed.py

@@ -1,6 +1,6 @@
 import torch
 
-from ..losses import BaseMetricLossFunction
+from ..losses import BaseMetricLossFunction, CrossBatchMemory
 from ..miners import BaseMiner
 from ..utils import common_functions as c_f
 from ..utils import loss_and_miner_utils as lmu
@@ -58,8 +58,14 @@ def get_indices_tuple(
 class DistributedLossWrapper(torch.nn.Module):
     def __init__(self, loss, efficient=False):
         super().__init__()
-        if not isinstance(loss, BaseMetricLossFunction):
-            raise TypeError("The input loss must extend BaseMetricLossFunction")
+        if not isinstance(loss, (BaseMetricLossFunction, CrossBatchMemory)):
+            raise TypeError(
+                "The input loss must extend BaseMetricLossFunction or CrossBatchMemory"
+            )
+        if isinstance(loss, CrossBatchMemory) and efficient:
+            raise ValueError(
+                "CrossBatchMemory with efficient=True is not currently supported"
+            )
         self.loss = loss
         self.efficient = efficient
 

src/pytorch_metric_learning/utils/inference.py

@@ -182,6 +182,10 @@ def __call__(
         c_f.LOGGER.info("embedding dimensionality is %d" % d)
         if self.reset_before:
             self.index = self.index_init_fn(d)
+        if self.index is None:
+            raise ValueError(
+                "self.index is None. It needs to be initialized before being used."
+            )
         distances, indices = try_gpu(
             self.index,
             query,

tests/utils/test_distributed.py

@@ -8,7 +8,8 @@
 import torch.optim as optim
 from torch.nn.parallel import DistributedDataParallel as DDP
 
-from pytorch_metric_learning import losses, miners
+from pytorch_metric_learning.losses import ContrastiveLoss, CrossBatchMemory
+from pytorch_metric_learning.miners import MultiSimilarityMiner
 from pytorch_metric_learning.utils import distributed
 
 from .. import TEST_DEVICE, TEST_DTYPES
@@ -52,6 +53,7 @@ def single_process_function(
     rank,
     world_size,
     lr,
+    iterations,
     model,
     inputs,
     labels,
@@ -83,20 +85,23 @@ def single_process_function(
         )
 
     optimizer = optim.SGD(ddp_mp_model.parameters(), lr=lr)
-    optimizer.zero_grad()
-    outputs = ddp_mp_model(inputs[rank].to(device))
-    indices_tuple = None
-    if miner_fn:
-        indices_tuple = miner_fn(outputs, labels[rank])
-    loss = loss_fn(outputs, labels[rank], indices_tuple)
-
-    dist.barrier()
-    loss.backward()
 
     original_model = original_model.to(device)
     assert not parameters_are_equal(original_model, ddp_mp_model.module)
-    dist.barrier()
-    optimizer.step()
+
+    for i in range(iterations):
+        optimizer.zero_grad()
+        outputs = ddp_mp_model(inputs[i][rank].to(device))
+        indices_tuple = None
+        if miner_fn:
+            indices_tuple = miner_fn(outputs, labels[i][rank])
+        loss = loss_fn(outputs, labels[i][rank], indices_tuple)
+
+        dist.barrier()
+        loss.backward()
+        dist.barrier()
+        optimizer.step()
+
     dist.barrier()
     assert parameters_are_equal(original_model, ddp_mp_model.module)
     dist.barrier()
@@ -113,7 +118,7 @@ def create_efficient_batch(x, i, batch_size):
 
 
 class TestDistributedLossWrapper(unittest.TestCase):
-    def loss_and_miner_tester(self, loss_class, miner_class, efficient):
+    def loss_and_miner_tester(self, loss_class, miner_class, efficient, xbm):
         torch.manual_seed(75210)
         if TEST_DEVICE == torch.device("cpu"):
             return
@@ -129,13 +134,7 @@ def loss_and_miner_tester(self, loss_class, miner_class, efficient):
             for world_size in range(2, max_world_size + 1):
                 batch_size = 20
                 lr = 1
-                inputs = [
-                    torch.randn(batch_size, 10).type(dtype) for _ in range(world_size)
-                ]
-                labels = [
-                    torch.randint(low=0, high=2, size=(batch_size,))
-                    for _ in range(world_size)
-                ]
+                iterations = 10
                 original_model = ToyMpModel().type(dtype)
                 model = ToyMpModel().type(dtype)
                 model.load_state_dict(original_model.state_dict())
@@ -144,6 +143,11 @@ def loss_and_miner_tester(self, loss_class, miner_class, efficient):
                 original_model = original_model.to(TEST_DEVICE)
                 original_loss_fn = loss_class()
                 loss_fn = loss_class()
+                if xbm:
+                    original_loss_fn = CrossBatchMemory(
+                        original_loss_fn, embedding_size=5
+                    )
+                    loss_fn = CrossBatchMemory(loss_fn, embedding_size=5)
 
                 if miner_class:
                     original_miner_fn = miner_class()
@@ -153,54 +157,68 @@ def loss_and_miner_tester(self, loss_class, miner_class, efficient):
                     miner_fn = None
 
                 optimizer = optim.SGD(original_model.parameters(), lr=lr)
-                optimizer.zero_grad()
-                all_inputs = torch.cat(inputs, dim=0).to(TEST_DEVICE)
-                all_labels = torch.cat(labels, dim=0).to(TEST_DEVICE)
-                all_outputs = original_model(all_inputs)
-                indices_tuple = None
-                if efficient:
-                    losses = []
-                    for i in range(len(inputs)):
-                        curr_emb, other_emb = create_efficient_batch(
-                            all_outputs, i, batch_size
-                        )
-                        curr_labels, other_labels = create_efficient_batch(
-                            all_labels, i, batch_size
-                        )
-                        if original_miner_fn:
-                            indices_tuple = distributed.get_indices_tuple(
+                inputs = [
+                    [torch.randn(batch_size, 10).type(dtype) for _ in range(world_size)]
+                    for _ in range(iterations)
+                ]
+                labels = [
+                    [
+                        torch.randint(low=0, high=2, size=(batch_size,))
+                        for _ in range(world_size)
+                    ]
+                    for _ in range(iterations)
+                ]
+
+                for aaa in range(iterations):
+                    optimizer.zero_grad()
+                    all_inputs = torch.cat(inputs[aaa], dim=0).to(TEST_DEVICE)
+                    all_labels = torch.cat(labels[aaa], dim=0).to(TEST_DEVICE)
+                    all_outputs = original_model(all_inputs)
+                    indices_tuple = None
+                    if efficient:
+                        losses = []
+                        for i in range(len(inputs[aaa])):
+                            curr_emb, other_emb = create_efficient_batch(
+                                all_outputs, i, batch_size
+                            )
+                            curr_labels, other_labels = create_efficient_batch(
+                                all_labels, i, batch_size
+                            )
+                            if original_miner_fn:
+                                indices_tuple = distributed.get_indices_tuple(
+                                    curr_labels,
+                                    other_labels,
+                                    TEST_DEVICE,
+                                    embeddings=curr_emb,
+                                    ref_emb=other_emb,
+                                    miner=original_miner_fn,
+                                )
+                            else:
+                                indices_tuple = distributed.get_indices_tuple(
+                                    curr_labels, other_labels, TEST_DEVICE
+                                )
+                            loss = original_loss_fn(
+                                curr_emb,
                                 curr_labels,
+                                indices_tuple,
+                                other_emb,
                                 other_labels,
-                                TEST_DEVICE,
-                                embeddings=curr_emb,
-                                ref_emb=other_emb,
-                                miner=original_miner_fn,
                             )
-                        else:
-                            indices_tuple = distributed.get_indices_tuple(
-                                curr_labels, other_labels, TEST_DEVICE
-                            )
-                        loss = original_loss_fn(
-                            curr_emb,
-                            curr_labels,
-                            indices_tuple,
-                            other_emb,
-                            other_labels,
-                        )
-                        losses.append(loss)
-                    loss = sum(losses)
-                else:
-                    if original_miner_fn:
-                        indices_tuple = original_miner_fn(all_outputs, all_labels)
-                    loss = original_loss_fn(all_outputs, all_labels, indices_tuple)
-                loss.backward()
-                optimizer.step()
+                            losses.append(loss)
+                        loss = sum(losses)
+                    else:
+                        if original_miner_fn:
+                            indices_tuple = original_miner_fn(all_outputs, all_labels)
+                        loss = original_loss_fn(all_outputs, all_labels, indices_tuple)
+                    loss.backward()
+                    optimizer.step()
 
                 mp.spawn(
                     single_process_function,
                     args=(
                         world_size,
                         lr,
+                        iterations,
                         model,
                         inputs,
                         labels,
@@ -215,19 +233,21 @@ def loss_and_miner_tester(self, loss_class, miner_class, efficient):
                 )
 
     def test_distributed_tuple_loss(self):
-        self.loss_and_miner_tester(losses.ContrastiveLoss, None, False)
+        for xbm in [False, True]:
+            self.loss_and_miner_tester(ContrastiveLoss, None, False, xbm)
 
     def test_distributed_tuple_loss_and_miner(self):
-        self.loss_and_miner_tester(
-            losses.ContrastiveLoss, miners.MultiSimilarityMiner, False
-        )
+        for xbm in [False, True]:
+            self.loss_and_miner_tester(
+                ContrastiveLoss, MultiSimilarityMiner, False, xbm
+            )
 
     def test_distributed_tuple_loss_efficient(self):
-        self.loss_and_miner_tester(losses.ContrastiveLoss, None, True)
+        self.loss_and_miner_tester(ContrastiveLoss, None, True, xbm=False)
 
     def test_distributed_tuple_loss_and_miner_efficient(self):
         self.loss_and_miner_tester(
-            losses.ContrastiveLoss, miners.MultiSimilarityMiner, True
+            ContrastiveLoss, MultiSimilarityMiner, True, xbm=False
         )
 
 

tests/utils/test_inference.py

@@ -34,7 +34,7 @@ class TestInference(unittest.TestCase):
     def setUpClass(cls):
         trunk = torchvision.models.resnet18(pretrained=True)
         trunk.fc = common_functions.Identity()
-        trunk = torch.nn.DataParallel(trunk.to(TEST_DEVICE))
+        trunk = trunk.to(TEST_DEVICE)
 
         cls.model = trunk
 
@@ -59,15 +59,17 @@ def tearDown(self):
         torch.cuda.empty_cache()
 
     def test_untrained_indexer(self):
-        inference_model = InferenceModel(trunk=self.model)
-        with self.assertRaises(RuntimeError):
-            inference_model.get_nearest_neighbors(self.dataset[0][0], k=10)
+        inference_model = InferenceModel(trunk=self.model, data_device=TEST_DEVICE)
+        with self.assertRaises(ValueError):
+            inference_model.get_nearest_neighbors(self.dataset[0][0].unsqueeze(0), k=10)
 
     def test_get_nearest_neighbors(self):
         test_filename = "test_inference.index"
         for indexer_input in [self.train_vectors, self.dataset]:
             for load_from_file in [False, True]:
-                inference_model = InferenceModel(trunk=self.model)
+                inference_model = InferenceModel(
+                    trunk=self.model, data_device=TEST_DEVICE
+                )
                 if load_from_file:
                     inference_model.load_knn_func(test_filename)
                 else:
@@ -79,15 +81,15 @@ def test_get_nearest_neighbors(self):
         os.remove(test_filename)
 
     def test_add_to_indexer(self):
-        inference_model = InferenceModel(trunk=self.model)
+        inference_model = InferenceModel(trunk=self.model, data_device=TEST_DEVICE)
         inference_model.knn_func.index = faiss.IndexFlatL2(512)
         inference_model.add_to_knn(self.dataset)
         self.helper_assertions(inference_model)
 
     def test_list_of_text(self):
         model = TextModel()
         dataset = TextDataset()
-        inference_model = InferenceModel(trunk=model)
+        inference_model = InferenceModel(trunk=model, data_device=TEST_DEVICE)
         inference_model.train_knn(dataset)
         inference_model.add_to_knn([["test1", "test2"], ["test3", "test4"]])
         result = inference_model.get_nearest_neighbors(["dog", "cat"], k=10)