Как использовать замаскированный выбор для такого рода тензоров

Question

Предположим, у меня есть тензор a и тензор b

import torch
a = torch.tensor([[[ 0.8856,  0.1411, -0.1856, -0.1425],
    [-0.0971,  0.1251,  0.1608, -0.1302],
    [-0.0901,  0.3215,  0.1763, -0.0412]], 
    [[ 0.8856,  0.1411, -0.1856, -0.1425],
    [-0.0971,  0.1251,  0.1608, -0.1302],
    [-0.0901,  0.3215,  0.1763, -0.0412]]])
b = torch.tensor([[0,
    2,
    1], 
    [0,
    2,
    1]])

Теперь я бы хотел выбрать индексы из тензора a, где значение тензора b равно не 0.

pred_masks = ( b != 0 )
c = torch.masked_select( a, (pred_masks == 1))

И, конечно, я получаю ожидаемую ошибку.

----> 1 c = torch.masked_select( a, (pred_masks == 1))

RuntimeError: The size of tensor a (4) must match the size of tensor b (3) at non-singleton dimension 2

Это вызвано вложенным списком, содержащим 4 элемента. Однако необходимо выбрать все значения вложенного списка по индексу x в тензоре a, соответствующем индексу x в тензоре b.

Буду благодарен за любую подсказку или ответ.

Answer 1

Я не совсем уверен, что вы хотите, чтобы форма вывода c. Так как ваша маска имеет форму (2,3) и имеет форму (2,3,4), вы хотите получить в качестве вывода тензор формы (n, 4), где n - это число элементов, которое истинно в ( 2,3) -маску?

Если да, то я бы предложил использовать маску в качестве индекса для первых двух измерений.

c = a[pred_masks,:]

Надеюсь, это немного поможет.