Python中多个张量的高效约简

Posted 2023-02-23

【中文标题】Python中多个张量的高效约简

【英文标题】：Efficient reduction of multiple tensors in Python

【发布时间】：2016-06-20 15:59:57

【问题描述】：

我在 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 中实现这部分，但我不确定应该使用什么算法。

【问题讨论】：

您能分享一下您的numpy.einsum 和numpy.tensordot 实现吗？

@Divakar：当然。 einsum 的实现就是 z = np.einsum('ijk,isl,jstm,ktn', v, a, w, x)。 tensordot 实现是` res = np.tensordot(v, a, (0,0)) res = np.tensordot(res, w, (0,0)) res = np.tensordot(res, x, ( 0,0)) `

请使用问题下方的“编辑”按钮将这些实现添加到问题中。

在您的 tensordot 中，您正在执行 z += ... 而没有提前实际初始化 z。你能澄清/纠正一下吗？

einsum 在所有列出的变量上构造一个迭代 (nditer)。我数了8，所以即使单个维度很小，产品迭代空间还是很大的。

【参考方案1】：

在部分使用np.tensordot，您可以像这样对事物进行矢量化 -

# 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

【讨论】：

谢谢！我尝试在我的程序中实现你的矢量化版本，它比我使用双 for 循环的版本快约 20%。

@Alessandro 感谢您报告性能改进！