PyTorch中一行与矩阵中其他行的逐元素批处理矩阵乘法



我有两个大小分别为(30, 24, 512)的矩阵,其中30是批次大小。让我们称它们为AB

现在我需要做的是:

对于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

最新更新