Replace squared_euclidean_distance_matrix
Summary:
`squared_euclidean_distance_matrix` is much slower to compute using `torch.norm` compared to doing matrix multiplication directly. See
I have personally measured a speedup by 70* , so I just reimplemented Natalia's code
@@ -16,7 +16,9 @@ def squared_euclidean_distance_matrix(pts1: torch.Tensor, pts2: torch.Tensor) ->
Tensor [M, N]: matrix of squared Euclidean distances; at index (m, n)
it contains || pts1[m] - pts2[n] ||^2
"""
- return torch.norm(pts1[:, None] - pts2, dim=2, p=2)
+ edm = torch.mm(-2 * pts1, pts2.t())
+ edm += (pts1 * pts1).sum(1, keepdim=True) + (pts2 * pts2).sum(1, keepdim=True).t()
+ return edm.contiguous()
def normalize_embeddings(embeddings: torch.Tensor, epsilon: float = 1e-6) -> torch.Tensor:
