CODEWITHSUNDEEP
pythonnumpypytorch
NicolaiF
NicolaiF

Reputation: 1343

Finding euclidean distance given an index array and a pytorch tensor

I got a pytorch tensor: 

Z = np.random.rand(100,2)
tZ = autograd.Variable(torch.cuda.FloatTensor(Z), requires_grad=True)

and an index array: 

idx = (np.array([0, 0, 0, 4, 3, 8], dtype="int64"),
       np.array([0, 1, 2, 3, 7, 4], dtype="int64"))

I need to find the distances of all the pairs of points in my tZ tensor using the idx array as indexes.

Right now I am doing it using numpy, but it would be nice if it could all be done using torch

dist = np.linalg.norm(tZ.cpu().data.numpy()[idx[0]]-tZ.cpu().data.numpy()[idx[1]], axis=1)

If anyone knows a way of doing it utilizing pytorch, to speed it up, it would be a great help!

Upvotes: 0

Views: 974

Answers (1)

benjaminplanche
benjaminplanche

Reputation: 15139

Using torch.index_select():

Z = np.random.rand(100,2)
tZ = autograd.Variable(torch.cuda.FloatTensor(Z), requires_grad=True)

idx = (np.array([0, 0, 0, 4, 3, 8], dtype="int64"),
       np.array([0, 1, 2, 3, 7, 4], dtype="int64"))

tZ_gathered = [torch.index_select(tZ, dim=0,
                                  index=torch.cuda.LongTensor(idx[i]))
                                  # note: you may have to wrap it in a Variable too
                                  # (thanks @rvd for the comment):
                                  # index = autograd.Variable(torch.cuda.LongTensor(idx[i])))
               for i in range(len(idx))]

print(tZ_gathered[0].shape)
# > torch.Size([6, 2])

dist = torch.norm(tZ_gathered[0] - tZ_gathered[1])

Upvotes: 1

Related Questions