Memory optimization for validation stage

oml/utils/misc_torch.py

@@ -98,6 +98,22 @@ def pairwise_dist(x1: Tensor, x2: Tensor, p: int = 2) -> Tensor:
     return cdist(x1=x1, x2=x2, p=p)
 
 
+def batched_knn(
+    query: torch.Tensor, gallery: torch.Tensor, k_neigh: int, batch_size: int
+) -> Tuple[torch.FloatTensor, torch.LongTensor]:
+    ids_neigh = torch.LongTensor(query.shape[0], k_neigh)
+    distances_neigh = torch.FloatTensor(query.shape[0], k_neigh)
+
+    for i in range(0, query.shape[0], batch_size):
+        distances_batch = pairwise_dist(x1=query[i : i + batch_size], x2=gallery)
+        distances_batch_neigh, ids_batch_neigh = torch.topk(distances_batch, k=k_neigh, largest=False, sorted=True)
+
+        ids_neigh[i : i + batch_size] = ids_batch_neigh
+        distances_neigh[i : i + batch_size] = distances_batch_neigh
+
+    return distances_neigh, ids_neigh
+
+
 def normalise(x: Tensor, p: int = 2) -> Tensor:
     """
     Args:
@@ -456,6 +472,7 @@ def _check_dimensions(self, n_components: int) -> None:
 __all__ = [
     "elementwise_dist",
     "pairwise_dist",
+    "batched_knn",
     "OnlineCalc",
     "AvgOnline",
     "SumOnline",

tests/test_oml/test_utils/test_misc_torch.py

@@ -7,8 +7,10 @@
 from oml.utils.misc_torch import (
     PCA,
     assign_2d,
+    batched_knn,
     drop_duplicates_by_ids,
     elementwise_dist,
+    pairwise_dist,
     take_2d,
 )
 
@@ -23,6 +25,22 @@ def test_elementwise_dist() -> None:
     assert torch.isclose(val_torch, torch.tensor(val_custom)).all()
 
 
+def test_batched_knn(n_generations: int = 5) -> None:
+    for _ in range(n_generations):
+        nq, ng, dim = 30, 50, 16
+        k, bs = 13, 9
+
+        query = torch.randn(nq, dim).float()
+        gallery = torch.randn(ng, dim).float()
+        distances_neigh, ids_neih = batched_knn(query, gallery, k_neigh=k, batch_size=bs)
+
+        distances = pairwise_dist(query, gallery)
+        distances_top_k_expected, ids_top_k_expected = torch.topk(distances, k=k, largest=False, sorted=True)
+
+        assert torch.isclose(distances_neigh, distances_top_k_expected).all()
+        assert (ids_neih == ids_top_k_expected).all()
+
+
 # fmt: off
 def test_take_2d() -> None:
     x = torch.tensor([