Pytorch batchwise matrix vector rowwise multiplication

user_14145452 0
Viewed 313

I am using pytorch.

If I have a matrix M of size (d1, d2) and a vector V of size d2, doing M*V gives me an output OUT of size (d1, d2), where each row of M has been multiplied by V.

I need to do the same thing batch-wise, where the matrix M is fixed and I have a batch of dB vectors.

In practice, given a tensor M of size (d1, d2) and a tensor V of size (dB, d2), I need to get as output OUT a tensor of size (dB, d1, d2), so that OUT[i] = M*V[i].

How can I efficiently get it with pytorch?

tensor pytorch python
2 Answers

This simple trick works for the problem:

M.unsqueeze(0) * V.unsqueeze(1)

This does multiplications of tensors having shapes (1, d1, d2) and (dB, 1, d2) and you get the desired output having shape (dB, d1, d2).

user_145272670

You can use Einstein Notation to achieve this:

torch.einsum('ij,bj->bij', M, V)
user_90676150
Related