使多维张量的numpy-einsum更快

Question

我有一些代码使用以下einsum：

y = np.einsum('wxyijk,ijkd->wxyd', x, f)

其中（例如）x的形状是（64，26，26，3，3，3），f的形状是（3，3，3，1），两者都具有dtype=float

%timeit np.einsum('wxyijk,ijkd->wxyd', x, f)
# 2.01 ms ± 55.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

这对于我的应用程序来说太慢了，这是时间关键。无论是使用GPU（通过CuPy）还是路径加速（通过opt-einsum）似乎都不会使速度更快。有没有办法让它在NumPy的本地速度更快，或者这是它将要达到的最快速度

Answer 1

在这种情况下，您可以使用optimize关键字，自己实现，或者使用tensordot。但是，第一个版本实际上应该做同样的事情（重塑->；点->；重塑）

您的实施

x=np.random.rand(64, 26, 26, 3, 3, 3)
f=np.random.rand(3, 3, 3, 1)
%timeit y = np.einsum('wxyijk,ijkd->wxyd', x, f)
#886 µs ± 3.16 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

使用optimize=“optimal”

%timeit y = np.einsum('wxyijk,ijkd->wxyd', x, f,optimize="optimal")
#275 µs ± 23.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

重塑和BLAS呼叫

这通常会导致与optimize="optimal"相当的性能，在这种情况下，不必要的阵列拷贝可能会导致性能下降

def contract(x,f):
    s1=x.shape
    x_=x.reshape(s1[0]*s1[1]*s1[2],s1[3]*s1[4]*s1[5])

    s2=f.shape
    f_=f.reshape(s2[0]*s2[1]*s2[2],s2[3])
    return np.dot(x_,f_).reshape(s1[0],s1[1],s1[2],s2[3])

%timeit contract(x,f)
#144 µs ± 3.09 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Tensordot

%timeit np.tensordot(x,f,axes=3)
#176 µs ± 4.32 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)