PyTorch L2-norm between 2 tensors of different shapes

1

I have 2 tensors in PyTorch:

a.shape, b.shape
# (torch.Size([1600, 2]), torch.Size([128, 2]))

I want to compute L2-norm distance between each of the 128 values in ‘b’ having 2-dim values from all 1600 values in ‘a’. Currently, I have an inefficient for loop to do it for each values in b as follows:

# Need to compute l2-norm squared dist b/w each b from a-
l2_dist_squared = list()

for bmu in bmu_locs:
    l2_dist_squared.append(torch.norm(input = a.to(torch.float32) - b, p = 2, dim = 1))

l2_dist_squared = torch.stack(l2_dist_squared)

# l2_dist_squared.shape
# torch.Size([128, 1600])

Is there a better way to do as a one liner?

2

Hi, you could use this functionality - torch.cdist — PyTorch 2.2 documentation which calculates the distance between each vector in ‘b’ to each vector in ‘a’.

E.g. torch.cdist(b, a, p=2) will give (128, 1600) shape, where each row represents the L2-norm distances between the corresponding value in ‘b’ and all values in ‘a’.

1 Like