使多维张量的 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)

张量点

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