Have Euclidean distance between two inputs populate a matrix within one loop

user_1545958 0
Viewed 236

I'm taking two tensors of dimensions (batch size, D1, D2, D3) and flattening them to (batch size, D1). I'm then trying to take the Euclidean distance between every row of one (train) tensor with every row of the second tensor (test). I'm having trouble understanding how to populate the distance combinations between tensors into each row element of the ith column.

  # Flatten
  train = x_train.view(num_train, x_train[1].view(1, -1).shape[1])
  test = x_test.view(num_test, x_test[1].view(1, -1).shape[1])

  # 1 Loop
  for i in range(num_test):
    dists[:,i] = torch.sqrt(torch.sum(torch.square(train-test[i])))

Doing it in one loop, every ith column is being populated with the same scalar value. I'm trying to make it such that

[1,i] = Euclidean distance between 1st image of train and 1st image of test ...
[2,i] = Euclidean distance between 1st image of train and 2nd image of test ...
[3,i] = Euclidean distance between 1st image of train and 3rd image of test
...
[(last element),i] = Euclidean distance between 1st image of train and (last element) image of test ...
[1,i+1] = Euclidean distance between 2nd image of train and 1st image of test
knn pytorch python
2 Answers

Okay, I was missing a very important parameter in torch.sum() that I did not know about that solves this issue. adding a 1, so that it looks like torch.sqrt(torch.sum(torch.square(train-test[i]), 1)) outputs to what I want.

user_15459580

Another way to do it is using torch.cdist:

torch.cdist(train.unsqueeze(0), test.unsqueeze(0), p=2)
user_59876980
Related