2 个 3D 张量之间张量点的维数

Question

我有一个关于 tensordot 操作的快速问题。我试图弄清楚是否有一种方法可以在两个张量之间执行张量点乘积以获得我想要的正确形状输出。其中一个张量是 B X L X D 维度，另一个是 B X 1 X D 维度，我想弄清楚是否有可能在最后得到 B X D 矩阵。

目前我正在遍历 B 维度并在 1 X D 和 D X L（转置 L X D）矩阵之间执行矩阵乘法并将它们堆叠以最终得到 B X L 矩阵。这显然不是最快的方法，因为循环可能很昂贵。是否可以通过执行快速张量点来获得所需的 B X D 形状输出？我似乎想不出摆脱其中一个 B 的方法。

任何见解或方向将不胜感激。

Answer 1

一个选项

就是使用torch.bmm()，它的作用正是(docs)。

它采用形状为 (b, n, m) 和 (b, m, p) 的张量和 returns 形状为 (b, n, p) 的批量矩阵乘法。

（我假设你得到了 B X L 的结果，因为 1 X D 和 D X L 的矩阵乘法的形状是 1 X L 而不是 1 X D）。

你的情况：

import torch
B, L, D = 32, 10, 512

a = torch.randn(B, 1, D)    #shape (B X 1 X D)
b = torch.randn(B, L, D)    #shape (B X L X D)

b = b.transpose(1,2)        #shape (B X D X L)

result = torch.bmm(a, b)

result = result.squeeze()
print(result.shape)
>>> torch.Size([32, 10])

或者

你可以使用torch.einsum()，我认为它更紧凑但可读性较差：

import torch
B, L, D = 32, 10, 512

a = torch.randn(B, 1, D)
b = torch.randn(B, L, D)

result = torch.einsum('abc, adc->ad', a, b)

print(result.shape)
>>> torch.Size([32, 10])

最后的挤压是为了让你的结果变成形状 (32, 10) 而不是形状 (32, 1, 10)。

Answer 2

我相信torch.einsum是执行张量求和的最直观方式：

>>> torch.einsum('bld,bed->bd', x, y)

其形状为 (B, D)。

明确表述，此处执行的操作等同于：

res = torch.zeros(B, D)
for b in range(B):
    for l in range(L):
        for d in range(D):
            res += x[b,l,d]*y[b,0,d]

_{实际上y上的第二个轴也循环了，但是范围只是[0]，因为y的第二个维度是一个单例。}

2 个 3D 张量之间张量点的维数

Dimension of tensordot between 2 3D tensors

python

pytorch

matrix-multiplication

dimension

tensordot

一个选项

或者