Revert to_device change (keep it for a future PR).

Author: Kevin Musgrave

src/pytorch_metric_learning/losses/manifold_loss.py

@@ -66,8 +66,7 @@ def compute_loss(self, embeddings, labels, indices_tuple, ref_emb, ref_labels):
             meta_classes = torch.cat((torch.arange(self.K), meta_classes))
             meta_classes = meta_classes[torch.randperm(N)]
 
-        loss_int = torch.zeros(1)
-        loss_int = c_f.to_device(loss_int, tensor=embeddings, dtype=embeddings.dtype)
+        loss_int = torch.zeros(1, device=embeddings.device, dtype=embeddings.dtype)
         embs_and_proxies = torch.cat([embeddings, self.proxies], dim=0)
 
         S = self.distance(embs_and_proxies, embs_and_proxies).clamp(0, np.inf)

src/pytorch_metric_learning/losses/p2s_grad_loss.py

@@ -44,8 +44,8 @@ def compute_loss(self, embeddings, labels, indices_tuple, ref_emb, ref_labels):
 
         self.weight.data = self.weight.data.renorm(2, 1, 1e-5).mul(1e5)
         dtype = embeddings.dtype
-        self.weight.data, labels = c_f.to_device(
-            (self.weight.data, labels), tensor=embeddings, dtype=(dtype, torch.long)
+        self.weight.data = c_f.to_device(
+            self.weight.data, tensor=embeddings, dtype=dtype
         )
 
         rtol = 1e-2 if dtype == torch.float16 else 1e-5

src/pytorch_metric_learning/samplers/hierarchical_sampler.py

@@ -74,12 +74,9 @@ def __len__(
     def reshuffle(self):
         batches = []
         for combinations in self.super_pairs:
-
             for b in range(self.batches_per_super_tuple):
-
                 batch = []
                 for slb in combinations:
-
                     sub_batch = []
                     all_classes = list(self.super_image_lists[slb].keys())
                     c_f.NUMPY_RANDOM.shuffle(all_classes)

src/pytorch_metric_learning/utils/common_functions.py

@@ -3,12 +3,10 @@
 import logging
 import os
 import re
-from typing import List, Tuple, Union
 
 import numpy as np
 import scipy.stats
 import torch
-from torch import nn
 
 LOGGER_NAME = "PML"
 LOGGER = logging.getLogger(LOGGER_NAME)
@@ -461,34 +459,13 @@ def to_dtype(x, tensor=None, dtype=None):
     return x
 
 
-def to_device(
-    x: Union[torch.Tensor, nn.Parameter, List, Tuple],
-    tensor=None,
-    device=None,
-    dtype: Union[torch.dtype, List, Tuple] = None,
-):
+def to_device(x, tensor=None, device=None, dtype=None):
     dv = device if device is not None else tensor.device
-    is_iterable = is_list_or_tuple(x)
-    if not is_iterable:
-        x = [x]
-
-    xd = x
-    if is_list_or_tuple(dtype):
-        if len(dtype) == len(x):
-            xd = [
-                to_dtype(x[i].to(dv), tensor=tensor, dtype=dtype[i])
-                for i in range(len(x))
-            ]
-        else:
-            raise RuntimeError(
-                f"The size of dtype was {len(dtype)}. It is only available 1 or the same of x"
-            )
-    elif dtype is not None:
-        xd = [to_dtype(xt.to(dv), tensor=tensor, dtype=dtype) for xt in x]
-
-    if len(xd) == 1:
-        xd = xd[0]
-    return xd
+    if x.device != dv:
+        x = x.to(dv)
+    if dtype is not None:
+        x = to_dtype(x, dtype=dtype)
+    return x
 
 
 def set_ref_emb(embeddings, labels, ref_emb, ref_labels):

tests/losses/test_manifold_loss.py

@@ -119,9 +119,7 @@ def forward(self, fvec, fLvec, fvecs_add=None):
             )
 
             for j in range(self.nb_proxy):
-
                 if fLvec[i] != j:
-
                     val1_context = (
                         self.d(torch.unsqueeze(A[i], 0), torch.unsqueeze(A_p[j], 0))
                         - dist_pos

tests/losses/test_multi_similarity_loss.py

@@ -12,7 +12,6 @@
 class TestMultiSimilarityLoss(unittest.TestCase):
     def test_multi_similarity_loss(self):
         for dtype in TEST_DTYPES:
-
             embedding_angles = [0, 20, 40, 60, 80]
             embeddings = torch.tensor(
                 [angle_to_coord(a) for a in embedding_angles],

tests/losses/test_p2s_grad_loss.py

@@ -137,9 +137,8 @@ def forward(self, input_score, target):
         # index (batch_size, class_num)
         with torch.no_grad():
             index = torch.zeros_like(input_score)
-            index, target = c_f.to_device(
-                (index, target), tensor=input_score, dtype=torch.long
-            )
+            index = c_f.to_device(index, tensor=input_score, dtype=torch.long)
+            target = c_f.to_device(target, tensor=input_score, dtype=torch.long)
             # index[i][target[i][j]] = 1
             index.scatter_(1, target.data.view(-1, 1), 1)
 
@@ -199,7 +198,6 @@ def test_p2s_grad_loss_with_paper_formula(self):
                 )
 
     def test_p2s_grad_loss_with_trusted_implementation(self):
-
         for dtype in TEST_DTYPES:
             embedding_angles = [0, 20, 40, 60, 80]
             embeddings = torch.tensor(

tests/losses/test_pnp_loss.py

@@ -31,7 +31,6 @@ def __init__(self, b, alpha, anneal, variant, bs, classes):
         self.mask = mask.unsqueeze(dim=0).repeat(self.batch_size, 1, 1)
 
     def forward(self, batch):
-
         dtype, device = batch.dtype, batch.device
         self.mask = self.mask.type(dtype).to(device)
         # compute the relevance scores via cosine similarity of the CNN-produced embedding vectors

tests/losses/test_proxy_anchor_loss.py

@@ -89,7 +89,6 @@ def test_proxyanchor_loss(self):
         margin = 0.5
 
         for use_autocast in [True, False]:
-
             if use_autocast:
                 cm = torch.cuda.amp.autocast()
             else:

tests/losses/test_triplet_margin_loss.py

@@ -14,7 +14,12 @@
 class TestTripletMarginLoss(unittest.TestCase):
     def test_triplet_margin_loss(self):
         for dtype in TEST_DTYPES:
-            embeddings = torch.randn(5, 32, requires_grad=True, dtype=dtype,).to(
+            embeddings = torch.randn(
+                5,
+                32,
+                requires_grad=True,
+                dtype=dtype,
+            ).to(
                 TEST_DEVICE
             )  # 2D embeddings
             embeddings = torch.nn.functional.normalize(embeddings)

tests/miners/test_batch_easy_hard_miner_labels.py

@@ -24,7 +24,7 @@ def test_labels(self):
                 DotProductSimilarity,
                 SNRDistance,
             ]:
-                for (pos_strategy, neg_strategy) in [
+                for pos_strategy, neg_strategy in [
                     ("easy", "easy"),
                     ("easy", "semihard"),
                     ("easy", "hard"),

tests/samplers/test_fixed_set_of_triplets.py

@@ -36,7 +36,7 @@ def test_fixed_set_of_triplets_with_batch_size(self):
                         dataset, batch_size=batch_size, sampler=sampler, drop_last=True
                     )
                     for _ in range(2):
-                        for (embeddings, curr_labels) in dataloader:
+                        for embeddings, curr_labels in dataloader:
                             a, p, n = miner(batch_of_fake_embeddings, curr_labels)
                             self.assertTrue(len(a) == batch_size // 3)
                             self.assertTrue(torch.all(curr_labels[a] == curr_labels[p]))

tests/samplers/test_m_per_class_sampler.py

@@ -68,7 +68,7 @@ def test_mperclass_sampler_with_batch_size(self):
                             drop_last=False,
                         )
                         for _ in range(2):
-                            for (_, curr_labels) in dataloader:
+                            for _, curr_labels in dataloader:
                                 unique_labels, counts = torch.unique(
                                     curr_labels, return_counts=True
                                 )

tests/trainers/test_key_checking.py

@@ -31,7 +31,6 @@ def test_metric_loss_only(self):
             TrainWithClassifier,
             TwoStreamMetricLoss,
         ]:
-
             model_dict = {"trunk": model}
             optimizer_dict = {"trunk_optimizer": None}
             loss_fn_dict = {"metric_loss": loss_fn}

tests/trainers/test_metric_loss_only.py

@@ -24,7 +24,6 @@
 
 class TestMetricLossOnly(unittest.TestCase):
     def test_metric_loss_only(self):
-
         cifar_resnet_folder = "temp_cifar_resnet_for_pytorch_metric_learning_test"
         dataset_folder = "temp_dataset_for_pytorch_metric_learning_test"
         model_folder = "temp_saved_models_for_pytorch_metric_learning_test"