Индексируйте тензор горелки с массивом

Question

У меня есть следующий тензор факела:

tensor([[-0.2,  0.3],
    [-0.5,  0.1],
    [-0.4,  0.2]])

и следующий массив numpy: (при необходимости я могу преобразовать его во что-то другое)

[1 0 1]

Я хочу получить следующий тензор:

tensor([0.3, -0.5, 0.2])

т.е. я хочу, чтобы массив numpy индексировал каждый подэлемент моего тензора. Желательно без использования al oop.

Заранее спасибо

Answer 1

Просто используйте диапазон (len (index)) для первого измерения.

import torch

a = torch.tensor([[-0.2,  0.3],
    [-0.5,  0.1],
    [-0.4,  0.2]])

c = [1, 0, 1]


b = a[range(3),c]

print(b)

Answer 2

Возможно, вы захотите использовать torch.gather - «Собирает значения вдоль оси, заданной dim.»

t = torch.tensor([[-0.2,  0.3],
    [-0.5,  0.1],
    [-0.4,  0.2]])
idxs = np.array([1,0,1])

idxs = torch.from_numpy(idxs).long().unsqueeze(1)  
# or   torch.from_numpy(idxs).long().view(-1,1)

t.gather(1, idxs)
tensor([[ 0.3000],
        [-0.5000],
        [ 0.2000]])

Здесь ваш индекс - массив numpy поэтому вы должны преобразовать его в LongTensor.