这个4D求和运算的张量是多少

Question

这里有一个简单的代码，可以将4D矩阵a与3D矩阵b:"批量相乘">

from functools import reduce
import numpy as np
from operator import mul
def einsum(a, b):
return np.einsum('ijkl,jkl->ikl', a, b)
def original(a, b):
s0, s1, s2, s3 = a.shape
c = np.empty((s0, s2, s3))
for j in range(s3):
for i in range(s2):
c[:, j, i] = np.dot(a[:, :, j, i], b[:, j, i])
return c
sz_a = (16, 4, 512, 512)
sz_b = (4, 512, 512)
a = np.random.random(reduce(mul, sz_a)).reshape(sz_a)
b = np.random.random(reduce(mul, sz_b)).reshape(sz_b)

定时：

%timeit original(a, b)
395 ms ± 2.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit einsum(a, b)
23.1 ms ± 191 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

我想测试tensordot的性能，看看它的比较情况，但我真的很难在这里使用它。如果有足够熟悉的人来指导我，我将不胜感激。非常感谢。

我最初的想法是：

np.tensordot(a, b, axes=((1),(0)))

但这给了我一个MemoryError，所以我认为这是不对的。。。

Answer 1

einsum与等效matmul的时间比较：

In [910]: timeit (a.transpose(2,3,0,1)@b[:,None].transpose(2,3,0,1)).transpose(2,3,0,1)[:
...: ,0]
90.5 ms ± 92.1 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [911]: timeit np.einsum('ijkl,jkl->ikl', a, b)
92.7 ms ± 2.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

时间已经很近了，我怀疑einsum优化实际上是在使用matmul。最初einsum使用自己编译的乘积和迭代，但最近随着最近的变化，它使用了各种方法，包括dot和matmul(如果合适的话(。

创建matmul是为了处理初始维度表示矩阵堆栈的情况。在您的问题中，最后两个维度是这个堆栈，dot作用于初始值。CCD_ 10是为了处理这种堆叠的点而创建的。CCD_ 11及其衍生物CCD_。