两个三维张量的点积

Question

两个三维张量的点积

4

我有两个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的形状为[批大小、向量数量、向量大小]。张量C应该表示来自A批中每个元素与来自B批中每个元素之间的点积，在所有不同向量之间进行计算。

希望这足够清晰，并期待您的回答！

- gorjan

3个回答

2

I think you can use einsum such as:

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

使用下标符号'ikm, jkm-> ijk'，您可以使用爱因斯坦约定指定要减少的维度。在此示例中，数组A和B的第三个维度均命名为'm'，将像向量上的点乘操作一样被减少。"最初的回答"。

- Ben.T

谢谢！你能把它翻译成矩阵乘法操作吗？ - gorjan

@gorjan，不太确定你想让我做什么？ - Ben.T

我在考虑，你的解决方案是否可以写成一个/多个矩阵乘法，而不是使用类似于einsum的方法。据我所知，einsum只是一种语法糖，它包装了一个或多个matmul的调用。 - gorjan

@gorjan 我明白你的意思，也许有一个简单的解决方案。但不幸的是，我似乎找不到一种不需要额外计算且与 Learning is a mess 中tensordot解决方案相同的方法。 - Ben.T

2

将数组转置为 kim,kmj->kij，然后使用 @。 - hpaulj

嘿@hpaulj，你能也发表一个答案吗？尽管已经有一个被接受的答案可以正常工作，但它在底层创建了一个巨大的张量来执行张量缩减，这使得在我的情况下使用它不可行。 - gorjan

-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 的张量。这可以通过取对角线来完成，就像将矩阵减少到其对角线条目的向量一样。

- Learning is a mess

请再试一次，第一个版本的轴索引有误。 - Learning is a mess

1

如果您能稍微解释一下会很有帮助。即使它有效，单行代码也不是很有用。 - Jovan Andonov

3

与上面的einsum相比，目前这种方法非常低效。您正在创建一个比C大N倍的数组，然后再取对角线。对于小数组，它至少比einsum慢10倍，对于大数组可能会慢几个数量级。 - Brenlla

@Brenlla 正确的。我刚刚对它们进行了基准测试，我的速度慢了10倍，而且临时内存占用更大。不过我不明白为什么会被踩。 - Learning is a mess

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- hpaulj · Accepted Answer

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会在最后两个维度上进行点积运算，并将前面的一个或多个维度作为批量维度。在你的情况下，'k'是批量维度，而'm'应该遵守“最后一个A和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，没有什么特别的。