Для заданного условия получите индексы значений в 2D-тензоре A, используйте их для индексации 3D-тензора B

Question

Для заданного 2D-тензора я хочу получить все индексы, где значение равно 1.Я ожидал, что смогу просто использовать torch.nonzero(a == 1).squeeze(), что вернет tensor([1, 3, 2]).Однако вместо этого torch.nonzero(a == 1) возвращает 2D-тензор (это нормально) с двумя значениями в строке (это не то, что я ожидал).Возвращенные индексы должны затем использоваться для индексации второго измерения (индекс 1) трехмерного тензора, снова возвращая двухмерный тензор.

import torch

a = torch.Tensor([[12, 1, 0, 0],
                  [4, 9, 21, 1],
                  [10, 2, 1, 0]])

b = torch.rand(3, 4, 8)

print('a_size', a.size())
# a_size torch.Size([3, 4])
print('b_size', b.size())
# b_size torch.Size([3, 4, 8])

idxs = torch.nonzero(a == 1)
print('idxs_size', idxs.size())
# idxs_size torch.Size([3, 2])

print(b.gather(1, idxs))

Очевидно, это не работает, что приводит к aRunTimeError:

RuntimeError: недопустимый аргумент 4: тензор индекса должен иметь те же размеры, что и тензор ввода в C: \ w \ 1 \ s \ windows \ pytorch \ aten \ src \ TH / generic / THTensorEvenMoreMath.cpp: 453

Кажется, что idxs - это не то, чего я ожидаю, и я не могу использовать его так, как я думал.idxs - это

tensor([[0, 1],
        [1, 3],
        [2, 2]])

, но, читая документацию , я не понимаю, почему я получаю обратно индексы строк в результирующем тензоре.Теперь я знаю, что могу получить правильные идентификаторы, нарезая idxs[:, 1], но, тем не менее, я не могу использовать эти значения в качестве индексов для трехмерного тензора, потому что возникает та же ошибка, что и раньше.Можно ли использовать одномерный тензор индексов для выбора элементов по данному измерению?

Answer 1

Вы можете просто нарезать их и передать в виде индексов, как в:

In [193]: idxs = torch.nonzero(a == 1)     
In [194]: c = b[idxs[:, 0], idxs[:, 1]]  

In [195]: c   
Out[195]: 
tensor([[0.3411, 0.3944, 0.8108, 0.3986, 0.3917, 0.1176, 0.6252, 0.4885],
        [0.5698, 0.3140, 0.6525, 0.7724, 0.3751, 0.3376, 0.5425, 0.1062],
        [0.7780, 0.4572, 0.5645, 0.5759, 0.5957, 0.2750, 0.6429, 0.1029]])

---

В качестве альтернативы, еще более простой и предпочтительный для меня подход - просто использовать torch.where(), а затем непосредственно индексировать в тензор b, как в:

In [196]: b[torch.where(a == 1)]  
Out[196]: 
tensor([[0.3411, 0.3944, 0.8108, 0.3986, 0.3917, 0.1176, 0.6252, 0.4885],
        [0.5698, 0.3140, 0.6525, 0.7724, 0.3751, 0.3376, 0.5425, 0.1062],
        [0.7780, 0.4572, 0.5645, 0.5759, 0.5957, 0.2750, 0.6429, 0.1029]])

Немного более подробного объяснения вышеуказанного подхода к использованию torch.where(): он работает на основе концепции расширенного индексирования .То есть, когда мы индексируем в тензор, используя кортеж объектов последовательности, таких как кортеж тензоров, кортеж списков, кортеж кортежей и т. Д.

# some input tensor
In [207]: a  
Out[207]: 
tensor([[12.,  1.,  0.,  0.],
        [ 4.,  9., 21.,  1.],
        [10.,  2.,  1.,  0.]])

Для базового среза нам понадобится кортеж целого числаиндексы:

   In [212]: a[(1, 2)] 
   Out[212]: tensor(21.)

Чтобы добиться того же с помощью расширенного индексирования, нам понадобится кортеж объектов последовательности:

# adv. indexing using a tuple of lists
In [213]: a[([1,], [2,])] 
Out[213]: tensor([21.])

# adv. indexing using a tuple of tuples
In [215]: a[((1,), (2,))]  
Out[215]: tensor([21.])

# adv. indexing using a tuple of tensors
In [214]: a[(torch.tensor([1,]), torch.tensor([2,]))] 
Out[214]: tensor([21.])

И размер возвращаемого тензора всегда будет на единицу меньшечем размерность входного тензора.

Answer 2

В дополнение к решению @ kmario23, вы все равно можете достичь тех же результатов, что и

b[torch.nonzero(a==1,as_tuple=True)]

Answer 3

Если предположить, что три измерения b равны batch_size x sequence_length x features (bxsx feats), ожидаемые результаты могут быть достигнуты следующим образом.

import torch

a = torch.Tensor([[12, 1, 0, 0],
                  [4, 9, 21, 1],
                  [10, 2, 1, 0]])

b = torch.rand(3, 4, 8)
print(b.size())
# b x s x feats
idxs = torch.nonzero(a == 1)[:, 1]
print(idxs.size())
# b
c = b[torch.arange(b.size(0)), idxs]
print(c.size())
# b x feats

Answer 4

import torch

a = torch.Tensor([[12, 1, 0, 0],
                  [4, 9, 21, 1],
                  [10, 2, 1, 0]])

b = torch.rand(3, 4, 8)

print('a_size', a.size())
# a_size torch.Size([3, 4])
print('b_size', b.size())
# b_size torch.Size([3, 4, 8])

#idxs = torch.nonzero(a == 1, as_tuple=True)
idxs = torch.nonzero(a == 1)
#print('idxs_size', idxs.size())

print(torch.index_select(b,1,idxs[:,1]))