Как векторизовать список матричных векторных умножений, используя pytorch / numpy

Question

Например, у меня есть список N B x H тензора (то есть N x B x H тензор) и список N векторов (т.е. N x B тензор).И я хочу умножить каждый B x H тензор в списке на соответствующий B размерный тензор, получив в результате N x H тензор.

Я знаю, как использовать один for-loop с PyTorch для реализации вычислений,но есть ли векторизация?(т. е. нет for-loop, просто используются операции PyTorch / numpy)

Answer 1

Вы можете достичь этого с помощью torch.bmm() и некоторых torch.squeeze() / torch.unsqueeze().

Мне лично больше нравятся более общие torch.einsum() (которые я считаю более читабельными):

import torch
import numpy as np

A = torch.from_numpy(np.array([[[1, 10, 100], [2, 20, 200], [3, 30, 300]],
                               [[4, 40, 400], [5, 50, 500], [6, 60, 600]]]))
B = torch.from_numpy(np.array([[ 1,  2,  3],
                               [-1, -2, -3]]))

AB = torch.einsum("nbh,nb->nh", (A, B))
print(AB)
# tensor([[   14,   140,  1400],
#         [  -32,  -320, -3200]])