在Python中有效减少多个张量

Question

我有四个多维张量v[i,j,k],a[i,s,l],w[j,s,t,m],x[k,t,n]在与NumPy,我试图来计算张量z[l,m,n]由下式给出:

z[l,m,n] = sum_{i,j,k,s,t} v[i,j,k] * a[i,s,l] * w[j,s,t,m] * x[k,t,n]

所有的张量相对较小(总共少于32k元素),但是我需要多次执行这个计算,所以我希望函数尽可能少地开销.

我尝试使用numpy.einsum这样实现它:

z = np.einsum('ijk,isl,jstm,ktn', v, a, w, x)

但它很慢.我也尝试了以下一系列numpy.tensordot调用:

z = np.zeros((a.shape[-1],w.shape[-1],x.shape[-1]))
for s in range(a.shape[1]):
  for t in range(x.shape[1]):
    res = np.tensordot(v, a[:,s,:], (0,0))
    res = np.tensordot(res, w[:,s,t,:], (0,0))
    z += np.tensordot(res, x[:,s,:], (0,0))

双在for循环中总结了s和t(包括s和t是非常小的,所以这是不是一个太大的问题).这样做效果更好,但仍然没有我想象的那么快.我认为这可能是因为tensordot在获取实际产品之前需要在内部执行的所有操作(例如,置换轴).

我想知道是否有更有效的方法在Numpy中实现这种操作.我也不介意在Cython中实现这个部分,但我不确定什么是正确的算法.

Answer 1

在部分中使用np.tensordot，您可以对事物进行矢量化，如下所示 -

\n\n

# Perform "np.einsum(\'ijk,isl->jksl\', v, a)"\np1 = np.tensordot(v,a,axes=([0],[0]))         # shape = jksl\n\n# Perform "np.einsum(\'jksl,jstm->kltm\', p1, w)"\np2 = np.tensordot(p1,w,axes=([0,2],[0,1]))    # shape = kltm\n\n# Perform "np.einsum(\'kltm,ktn->lmn\', p2, w)"\nz = np.tensordot(p2,x,axes=([0,2],[0,1]))     # shape = lmn\n

\n\n

运行时测试并验证输出 -

\n\n

In [15]: def einsum_based(v, a, w, x):\n    ...:     return np.einsum(\'ijk,isl,jstm,ktn\', v, a, w, x) # (l,m,n)\n    ...: \n    ...: def vectorized_tdot(v, a, w, x):\n    ...:     p1 = np.tensordot(v,a,axes=([0],[0]))        # shape = jksl\n    ...:     p2 = np.tensordot(p1,w,axes=([0,2],[0,1]))   # shape = kltm\n    ...:     return np.tensordot(p2,x,axes=([0,2],[0,1])) # shape = lmn\n    ...: \n

\n\n

情况1 ：

\n\n

In [16]: # Input params\n    ...: i,j,k,l,m,n = 10,10,10,10,10,10\n    ...: s,t = 3,3 # As problem states : "both s and t are very small".\n    ...: \n    ...: # Input arrays\n    ...: v = np.random.rand(i,j,k)\n    ...: a = np.random.rand(i,s,l)\n    ...: w = np.random.rand(j,s,t,m)\n    ...: x = np.random.rand(k,t,n)\n    ...: \n\nIn [17]: np.allclose(einsum_based(v, a, w, x),vectorized_tdot(v, a, w, x))\nOut[17]: True\n\nIn [18]: %timeit einsum_based(v,a,w,x)\n10 loops, best of 3: 129 ms per loop\n\nIn [19]: %timeit vectorized_tdot(v,a,w,x)\n1000 loops, best of 3: 397 \xc2\xb5s per loop\n

\n\n

情况#2（更大的数据量）：

\n\n

In [20]: # Input params\n    ...: i,j,k,l,m,n = 15,15,15,15,15,15\n    ...: s,t = 3,3 # As problem states : "both s and t are very small".\n    ...: \n    ...: # Input arrays\n    ...: v = np.random.rand(i,j,k)\n    ...: a = np.random.rand(i,s,l)\n    ...: w = np.random.rand(j,s,t,m)\n    ...: x = np.random.rand(k,t,n)\n    ...: \n\nIn [21]: np.allclose(einsum_based(v, a, w, x),vectorized_tdot(v, a, w, x))\nOut[21]: True\n\nIn [22]: %timeit einsum_based(v,a,w,x)\n1 loops, best of 3: 1.35 s per loop\n\nIn [23]: %timeit vectorized_tdot(v,a,w,x)\n1000 loops, best of 3: 1.52 ms per loop\n

\n