python - Pytorch 跨不同数组的行进行点积

Question

我正在尝试编写类似于变压器论文中的位置编码的代码。为此，我需要执行以下操作:

对于以下三个矩阵，我想在行级别将它们连接起来(即每个矩阵的第一行堆叠在一起，第二行堆叠在一起，等等)，然后在每个矩阵及其转置之间应用点积，并且最后，将它们压平并叠在一起。我将在以下示例中阐明这一点:

x = torch.tensor([[1,1,1,1],
                  [2,2,2,2],
                  [3,3,3,3]])
y = torch.tensor([[0,0,0,0],
                  [0,0,0,0],
                  [0,0,0,0]])
z = torch.tensor([[4,4,4,4],
                  [5,5,5,5],
                  [6,6,6,6]])

concat = torch.cat([x, y, z], dim=-1).view(-1, x.shape[-1])
print(concat)

tensor([[1, 1, 1, 1],
        [0, 0, 0, 0],
        [4, 4, 4, 4],
        [2, 2, 2, 2],
        [0, 0, 0, 0],
        [5, 5, 5, 5],
        [3, 3, 3, 3],
        [0, 0, 0, 0],
        [6, 6, 6, 6]])

# Here I get each three rows together, and then apply dot product, flatten, and stack them.
concat = torch.stack([
            torch.flatten(
                torch.matmul(
                    concat[i:i+3, :], # 3 is the number of tensors (x,y,z)
                    torch.transpose(concat[i:i+3, :], 0, 1))
                )
            for i in range(0, concat.shape[0], 3)
            ])

print(concat)

tensor([[  4,   0,  16,   0,   0,   0,  16,   0,  64],
        [ 16,   0,  40,   0,   0,   0,  40,   0, 100],
        [ 36,   0,  72,   0,   0,   0,  72,   0, 144]])

终于，我得到了我想要的最终矩阵。我的问题是，有没有办法像我在最后一步中那样使用循环来实现这一目标？我希望一切都在张量中。

Answer 1

torch.einsum 使matmul您想要的轴变得更容易。

c = torch.concat([x, y, z], dim=-1).reshape(-1, *x.shape)
torch.einsum('ijl,ikl->ikj', c, c).reshape(3, -1)

输出

tensor([[  4,   0,  16,   0,   0,   0,  16,   0,  64],
        [ 16,   0,  40,   0,   0,   0,  40,   0, 100],
        [ 36,   0,  72,   0,   0,   0,  72,   0, 144]])