Initial commit

2026-07-13 01:00:43 +08:00 · 2017-12-26 18:43:17 -05:00
commit dc4e2b0db3
42 changed files with 3486 additions and 0 deletions
@@ -0,0 +1,76 @@
+import torch
+from enum import Enum
+
+PDist2Order = Enum('PDist2Order', 'd_first d_second')
+
+
+def pdist2(X: torch.Tensor,
+           Z: torch.Tensor = None,
+           order: PDist2Order = PDist2Order.d_second) -> torch.Tensor:
+    r""" Calculates the pairwise distance between X and Z
+
+    D[b, i, j] = l2 distance X[b, i] and Z[b, j]
+
+    Parameters
+    ---------
+    X : torch.Tensor
+        X is a (B, N, d) tensor.  There are B batches, and N vectors of dimension d
+    Z: torch.Tensor
+        Z is a (B, M, d) tensor.  If Z is None, then Z = X
+
+    Returns
+    -------
+    torch.Tensor
+        Distance matrix is size (B, N, M)
+    """
+
+    if order == PDist2Order.d_second:
+        if X.dim() == 2:
+            X = X.unsqueeze(0)
+        if Z is None:
+            Z = X
+            G = X @ Z.transpose(-2, -1)
+            S = (X * X).sum(-1, keepdim=True)
+            R = S.transpose(-2, -1)
+        else:
+            if Z.dim() == 2:
+                Z = Z.unsqueeze(0)
+            G = X @ Z.transpose(-2, -1)
+            S = (X * X).sum(-1, keepdim=True)
+            R = (Z * Z).sum(-1, keepdim=True).transpose(-2, -1)
+    else:
+        if X.dim() == 2:
+            X = X.unsqueeze(0)
+        if Z is None:
+            Z = X
+            G = X.transpose(-2, -1) @ Z
+            R = (X * X).sum(-2, keepdim=True)
+            S = R.transpose(-2, -1)
+        else:
+            if Z.dim() == 2:
+                Z = Z.unsqueeze(0)
+            G = X.transpose(-2, -1) @ Z
+            S = (X * X).sum(-2, keepdim=True).transpose(-2, -1)
+            R = (Z * Z).sum(-2, keepdim=True)
+
+    return torch.abs(R + S - 2 * G).squeeze(0)
+
+
+def pdist2_slow(X, Z=None):
+    if Z is None: Z = X
+    D = torch.zeros(X.size(0), X.size(2), Z.size(2))
+
+    for b in range(D.size(0)):
+        for i in range(D.size(1)):
+            for j in range(D.size(2)):
+                D[b, i, j] = torch.dist(X[b, :, i], Z[b, :, j])
+    return D
+
+
+if __name__ == "__main__":
+    X = torch.randn(2, 3, 5)
+    Z = torch.randn(2, 3, 3)
+
+    print(pdist2(X, order=PDist2Order.d_first))
+    print(pdist2_slow(X))
+    print(torch.dist(pdist2(X, order=PDist2Order.d_first), pdist2_slow(X)))