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'，您可以指定使用爱因斯坦惯例缩小的尺寸。数组A和B的第三个维（这里命名为'm'都将像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给出了一个不同的轴而不是数学顺序）。这可以通过采用对角线来完成，就像将矩阵缩小为对角线项的向量时一样。

两个3D张量之间的点积

3 个答案: