Нахождение евклидова расстояния по массиву индексов и тензору pytorch - PullRequest
Новая
       13

Нахождение евклидова расстояния по массиву индексов и тензору pytorch

0 голосов
/ 02 июня 2018

У меня есть тензор Pytorch: 

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, используямассив idx в качестве индексов.

Прямо сейчас я делаю это с помощью numpy, но было бы неплохо, если бы все это можно было сделать с помощью факела 

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

Если кто-нибудь знает способ сделать это с помощью pytorch, чтобы ускоритьэто было бы здорово помочь!

1 Ответ

0 голосов
/ 02 июня 2018

Использование 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])
...