两个 3D 张量之间的点积

Question

我有两个 3D 张量，张量 A 的形状为 [B,N,S]，张量 B 的形状也为 [B,N,S]。我想要得到的是第三张量C，我希望它具有[B,B,N] 形状，其中元素C[i,j,k] = np.dot(A[i,k,:], B[j,k,:]。我也想实现这是一种矢量化的方式。

更多信息：A 和 B 两个张量的形状为 [Batch_size, Num_vectors, Vector_size]。张量C，应该表示来自A 的批次中的每个元素与来自B 的批次中的每个元素之间的点积，在所有不同的向量之间。

希望够清楚，期待你的解答！

Answer 1

In [331]: A=np.random.rand(100,200,300)                                                              
In [332]: B=A

建议的einsum，直接从

C[i,j,k] = np.dot(A[i,k,:], B[j,k,:] 

表达式：

In [333]: np.einsum( 'ikm, jkm-> ijk', A, B).shape                                                   
Out[333]: (100, 100, 200)
In [334]: timeit np.einsum( 'ikm, jkm-> ijk', A, B).shape                                            
800 ms ± 25.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

matmul 在最后 2 个维度上执行 dot，并将领先的维度视为批处理。在您的情况下，“k”是批次维度，“m”是应该遵守last A and 2nd to the last of B 规则的维度。所以重写ikm,jkm... 以适应，并相应地转置A 和B：

In [335]: np.einsum('kim,kmj->kij', A.transpose(1,0,2), B.transpose(1,2,0)).shape                     
Out[335]: (200, 100, 100)
In [336]: timeit np.einsum('kim,kmj->kij',A.transpose(1,0,2), B.transpose(1,2,0)).shape              
774 ms ± 22.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

性能差别不大。但是现在使用matmul：

In [337]: (A.transpose(1,0,2)@B.transpose(1,2,0)).transpose(1,2,0).shape                             
Out[337]: (100, 100, 200)
In [338]: timeit (A.transpose(1,0,2)@B.transpose(1,2,0)).transpose(1,2,0).shape                      
64.4 ms ± 1.17 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

并验证值是否匹配（但通常情况下，如果形状匹配，值也会匹配）。

In [339]: np.allclose((A.transpose(1,0,2)@B.transpose(1,2,0)).transpose(1,2,0),np.einsum( 'ikm, jkm->
     ...:  ijk', A, B))                                                                              
Out[339]: True

我不会尝试测量内存使用情况，但时间改进表明它也更好。

在某些情况下，einsum 被优化为使用matmul。这里似乎不是这样，尽管我们可以使用它的参数。我有点惊讶matmul 做得这么好。

===

我隐约记得另一个关于matmul 在两个数组相同时采取捷径的 SO，A@A。我在这些测试中使用了B=A。

In [350]: timeit (A.transpose(1,0,2)@B.transpose(1,2,0)).transpose(1,2,0).shape                      
60.6 ms ± 1.17 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [352]: B2=np.random.rand(100,200,300)                                                             
In [353]: timeit (A.transpose(1,0,2)@B2.transpose(1,2,0)).transpose(1,2,0).shape                     
97.4 ms ± 164 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

但这只会产生轻微的影响。

In [356]: np.__version__                                                                             
Out[356]: '1.16.4'

我的 BLAS 等是标准的 Linux，没什么特别的。

Answer 2

我想你可以使用einsum 如：

np.einsum( 'ikm, jkm-> ijk', A, B)

使用下标'ikm, jkm-> ijk'，您可以指定使用爱因斯坦约定减少的维度。此处名为 'm' 的数组 A 和 B 的第三维将被缩减，就像 dot 操作对向量所做的那样。

Answer 3

试试：

C = np.diagonal( np.tensordot(A,B, axes=(2,2)), axis1=1, axis2=3)

来自https://docs.scipy.org/doc/numpy/reference/generated/numpy.tensordot.html#numpy.tensordot

解释

解决方案是两个操作的组合。首先是 A 和 B 在其第三轴上的张量积，如您所愿。这会输出一个秩为 4 的张量，您希望通过在轴 1 和轴 3 上取相等的索引来将其减少到一个秩为 3 的张量（您的符号中的 k，请注意，tensordot 给出的轴顺序与您的数学不同）。这可以通过取对角线来完成，就像将矩阵简化为其对角线元素的向量一样。