Python 使多维张量的numpy-einsum更快

我有一些代码使用以下

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中，有没有任何方法可以让它在本机上运行得更快，或者这是最快的速度？

在本例中，您可以使用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”

重塑和BLAS呼叫

%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)

这通常会导致与

optimize=“optimimal”

相当的性能，在这种情况下，不必要的阵列拷贝可能会导致性能下降

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)

---

您可以将数组的形状改为（64,26,26,27）和（27,1），并使用

@

。但在我的测试中，这并不能提高速度。对此不确定，但100个循环可能不足以将字符串的初始解析分摊到一组操作中

tensordot

会重新成形（并在需要时转置）将计算减少到

contract

中的

np.dot

。由于开发人员尝试了不同的策略，在

einsum

中的

optimize

的效果在numpy版本上发生了变化。在我的版本（1.18）中，它没有帮助。@hpaulj我还有1.18.1。使用tensordot对您的系统有影响吗？也许它是这个非常特殊情况下的BLAS后端（使用点->gemm进行矩阵向量乘法）。我有MKLSTRANGLY，在我的系统上，naive版本需要2毫秒，contract和tensordot都需要4.92毫秒，而“optimal”需要惊人的5.42毫秒！有什么设置需要更改吗？@diehardthethryhard我在CoreI5-8500（双通道RAM）上计时。我也在低端设备上试用过（带有core-m3和单通道RAM的平板电脑）。每种方法在0.88～0.95ms之间。你的计时对那个硬件来说太慢了。我想最好是安装一个干净的水蟒并在那里进行测试。