一行与矩阵中每一行的逐元素批处理矩阵乘法，单位为PyTorch

我有两个大小分别为

（30,24,512）

的矩阵，其中

30

是批量大小。让我们称它们为

A

和

B

现在我需要做的是：

对于

A

中的每一批，我想计算

A

单个批中每一行与

B

单个批中每一行的元素级批矩阵乘法，并求和。 换句话说，对于每个批次，我在左侧（

a

）和右侧（

B

）都有一个（24512）矩阵。现在，对于

A

中的每一行

（1*512）

，我想计算该行与B中的每一行

（24*512）

的元素矩阵乘法，并逐一求和。 因此，此操作将产生

（1*512）

大小的向量。 对A中的每一行执行，将生成一个

（24*512）

大小的矩阵，然后对每一批执行，将生成

（30,24,512）

大小的矩阵，这是我输出的首选维度

我该怎么做呢？

---

我不想使用for循环，因为那样效率会很低。出于同样的原因，我也不想使用

repeat

方法。

我找到了自己问题的答案。我没有计算每行的元素乘积，而是计算行的和，然后取元素乘积，这使得问题（和解决方案）更简单

import torch

a = torch.randn(30, 24, 512)
b = torch.randn(30, 24, 512)

# Step 1: Summing b along dimension 1 and unsqueezing
# Dimension of b_sum is 30*1*512
b_sum = torch.sum(b, dim=1).unsqueeze(1)

# Step 2: Element-wise product 
result = a * b_sum

---

请用绿色勾号标出您的答案，以供未来读者参考，谢谢参与！你可以阅读@SzymonMaszke哦，是的，我之前确实尝试过，但我得到一个错误，说我只能明天接受我的答案。