Как я могу получить один массив, который будет возвращать только маскированные значения, определенные другим массивом с Numpy / PyTorch?

Question

У меня есть mask, который имеет форму: [64, 2895] и массив pred, который имеет форму [64, 2895, 161].

mask является двоичным с только 0 с и 1 с. То, что я хочу сделать, это уменьшить pred, чтобы он поддерживал 64 пакетов, и вдоль 2895, где есть 1 в mask для каждого пакета, вернуть соответствующий pred.

Итак, в качестве упрощенного примера, если:

mask = [[1, 0, 0],
        [1, 1, 0],
        [0, 0, 1]]
pred = [[[0.12, 0.23, 0.45, 0.56, 0.57],
         [0.91, 0.98, 0.97, 0.96, 0.95],
         [0.24, 0.46, 0.68, 0.80, 0.15]],

        [[1.12, 1.23, 1.45, 1.56, 1.57],
         [1.91, 1.98, 1.97, 1.96, 1.95],
         [1.24, 1.46, 1.68, 1.80, 1.15]],

        [[2.12, 2.23, 2.45, 2.56, 2.57],
         [2.91, 2.98, 2.97, 2.96, 2.95],
         [2.24, 2.46, 2.68, 2.80, 2.15]]]

То, что я хочу, это:

[[[0.12, 0.23, 0.45, 0.56, 0.57]],

 [[1.12, 1.23, 1.45, 1.56, 1.57],
  [1.91, 1.98, 1.97, 1.96, 1.95]],

 [[2.24, 2.46, 2.68, 2.80, 2.15]]]

Я понимаю, что существуют разные измерения, я надеюсь, что это возможно. Если нет, то заполните недостающие размеры 0. Либо numpy или pytorch будет полезно. Спасибо.

Answer 1

Если вы хотите векторизованное вычисление, тогда другое измерение кажется невозможным, но это даст вам тот, у которого запись в маске заполнена 0:

# pred: torch.size([64, 2895, 161])
# mask: torch.size([64, 2895])

result = pred * mask[:, :, None] 
# extend mask with another dimension so now it can do entry-wise multiplication

и result - это именно то, что вы хотите