pytorch - 在 PyTorch 中计算两个相同大小的方阵的逐行点积的有效方法

Question

假设我有两个大小相同的方阵 A、B

A = torch.tensor([[1, 2], [3, 4]])
B = torch.tensor([[1, 1], [1, 1]])

我想要一个由逐行点积组成的结果张量，比如说

tensor([3, 7])  # i.e. (1*1 + 2*1, 3*1 + 4*1) 

在 PyTorch 中实现这一目标的有效方法是什么？

标签： pytorch

Answer 1

正如你所说，你可以使用torch.bmm，但你首先需要广播你的输入：

>>> torch.bmm(A[..., None, :], B[..., None])
tensor([[[3]],

        [[7]]])

或者，您可以使用torch.einsum：

>>> torch.einsum('ij,ij->i', A, B)
tensor([3, 7])