多次循环 np.einsum ...有更快的方法吗？答案

【问题标题】：Looping over np.einsum many times... Is there a faster way?多次循环 np.einsum ...有更快的方法吗？
【发布时间】：2020-11-04 03:20:54
【问题描述】：

我有一个似然函数，我正在尝试使用 MCMC 进行采样。我在对数似然本身中没有使用 for 循环，但我确实调用了一次 np.einsum()。

这是我当前代码的示例：

A = np.random.rand(4,50,60,200) # Random NDarray
B = np.random.rand(200,1000,4)  # Random NDarray
out = np.einsum('ijkl,lui->jkui', A, B, optimize="optimal")

输出out 的尺寸为 (50,60,1000,4)。这个计算有点太慢了，无法进行有效的 MCMC 采样（我的机器上大约 4 秒），有什么办法可以加快速度吗？ 一个有用的信息是，对于对数似然函数的每次调用，虽然数组 A 和 B 中的实际值在变化，但每个数组的维度保持不变。我想这是可能有助于加快速度，因为总是将相同的元素相乘。

【问题讨论】：

标签： python performance numpy numpy-einsum log-likelihood

【解决方案1】：

其中一个轴在A（第一个）和B（最后一个）中保持对齐，并且也保持在输出中（最后一个），并且是4 的一个非常小的循环数。所以，我们可以简单地用 np.tensordot 循环那个，以减少张量。在处理如此大的数据集时，4x 较少的内存拥塞的好处可能会克服 4x 循环，因为每次迭代的计算量也较少 4x。

因此，tensordot 的解决方案将是 -

def func1(A, B):
    out = np.empty(A.shape[1:3] + B.shape[1:])
    for i in range(len(A)):
        out[...,i] = np.tensordot(A[i], B[...,i],axes=(-1,0))
    return out

时间安排 -

In [70]: A = np.random.rand(4,50,60,200) # Random NDarray
    ...: B = np.random.rand(200,1000,4)  # Random NDarray
    ...: out = np.einsum('ijkl,lui->jkui', A, B, optimize="optimal")

# Einsum solution without optimize    
In [71]: %timeit np.einsum('ijkl,lui->jkui', A, B)
2.89 s ± 109 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# Einsum solution with optimize    
In [72]: %timeit np.einsum('ijkl,lui->jkui', A, B, optimize="optimal")
2.79 s ± 9.31 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)    

# @Paul Panzer's soln
In [74]: %timeit np.stack([np.tensordot(a,b,1) for a,b in zip(A,B.transpose(2,0,1))],-1)
183 ms ± 6.08 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [73]: %timeit func1(A,B)
158 ms ± 3.35 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

只是为了重申内存拥塞和计算需求的重要性，假设我们也想对长度的最后一个轴进行求和减少4，那么我们将看到optimal 的时序有明显差异版本-

In [78]: %timeit np.einsum('ijkl,lui->jkui', A, B, optimize="optimal")
2.76 s ± 9.36 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [79]: %timeit np.einsum('ijkl,lui->jku', A, B, optimize="optimal")
93.8 ms ± 3.3 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

所以，在这种情况下，最好选择einsum。

特定于给定问题

鉴于A 和B 的维度保持不变，使用out = np.empty(A.shape[1:3] + B.shape[1:]) 的数组初始化可以一次性完成，并使用建议的循环遍历对数似然函数的每次调用过度使用tensordot 并更新输出out。

【讨论】：

@Jsn 如果您是 tensordot 的新手，这可能会有所帮助 - stackoverflow.com/questions/41870228 提供一些示例。
对内存拥塞的见解很有启发意义。我不知道效果如此剧烈。
内存拥塞示例更多的是 einsum 与 BLAS 的问题。 78 将重组张量并执行 TDOT，79 无法调用 TDOT，因此它使用 einsum 循环。

【解决方案2】：

即使在小循环中使用 tensordot 也快 10 倍以上：

timeit(lambda:np.einsum('ijkl,lui->jkui', A, B, optimize="optimal"),number=5)/5
# 3.052245747600682
timeit(lambda:np.stack([np.tensordot(a,b,1) for a,b in zip(A,B.transpose(2,0,1))],-1),number=10)/10
# 0.23842503569903784

out_td = np.stack([np.tensordot(a,b,1) for a,b in zip(A,B.transpose(2,0,1))],-1)
out_es = np.einsum('ijkl,lui->jkui', A, B, optimize="optimal")
np.allclose(out_td,out_es)
# True

【讨论】：