两个 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

尝试：

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 上采用相等的索引（您的 k 在您的符号中，请注意 tensordot 给出了不同的轴来将其减少为 3 阶张量顺序比你的数学）。这可以通过采用对角线来完成，就像将矩阵简化为其对角线元素的向量时所做的那样。

Answer 2

我认为你可以使用 einsum 例如：

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

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

Answer 3

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 在最后两个维度上执行 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，没什么特别的。

两个 3D 张量之间的点积

Dot product between two 3D tensors

python

numpy

tensorflow

tensordot