Python中多重张量的有效归约

Question

我在Numpy中有四个多维张量v[i,j,k]，a[i,s,l]，w[j,s,t,m]，x[k,t,n]，我试图计算张量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

部分使用^{}，你可以像这样矢量化-

# Perform "np.einsum('ijk,isl->jksl', v, a)"
p1 = np.tensordot(v,a,axes=([0],[0]))         # shape = jksl

# Perform "np.einsum('jksl,jstm->kltm', p1, w)"
p2 = np.tensordot(p1,w,axes=([0,2],[0,1]))    # shape = kltm

# Perform "np.einsum('kltm,ktn->lmn', p2, w)"
z = np.tensordot(p2,x,axes=([0,2],[0,1]))     # shape = lmn

运行时测试和验证输出-

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

案例1：

In [16]: # Input params
    ...: i,j,k,l,m,n = 10,10,10,10,10,10
    ...: s,t = 3,3 # As problem states : "both s and t are very small".
    ...: 
    ...: # Input arrays
    ...: v = np.random.rand(i,j,k)
    ...: a = np.random.rand(i,s,l)
    ...: w = np.random.rand(j,s,t,m)
    ...: x = np.random.rand(k,t,n)
    ...: 

In [17]: np.allclose(einsum_based(v, a, w, x),vectorized_tdot(v, a, w, x))
Out[17]: True

In [18]: %timeit einsum_based(v,a,w,x)
10 loops, best of 3: 129 ms per loop

In [19]: %timeit vectorized_tdot(v,a,w,x)
1000 loops, best of 3: 397 µs per loop

案例2（更大的数据量）：

In [20]: # Input params
    ...: i,j,k,l,m,n = 15,15,15,15,15,15
    ...: s,t = 3,3 # As problem states : "both s and t are very small".
    ...: 
    ...: # Input arrays
    ...: v = np.random.rand(i,j,k)
    ...: a = np.random.rand(i,s,l)
    ...: w = np.random.rand(j,s,t,m)
    ...: x = np.random.rand(k,t,n)
    ...: 

In [21]: np.allclose(einsum_based(v, a, w, x),vectorized_tdot(v, a, w, x))
Out[21]: True

In [22]: %timeit einsum_based(v,a,w,x)
1 loops, best of 3: 1.35 s per loop

In [23]: %timeit vectorized_tdot(v,a,w,x)
1000 loops, best of 3: 1.52 ms per loop