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

Question

我有一个似然函数，我正在尝试使用 MCMC 对其进行采样。我在对数似然本身中没有使用 for 循环，但我确实调用了np.einsum()一次。这是我当前代码的示例：A = np.random.rand(4,50,60,200) # Random NDarrayB = np.random.rand(200,1000,4)  # Random NDarrayout = np.einsum('ijkl,lui->jkui', A, B, optimize="optimal")输出out具有维度 (50,60,1000,4)。这个计算有点太慢，无法进行有效的 MCMC 采样（在我的机器上大约 4 秒），有什么办法可以加快速度吗？一条有用的信息是，对于对数似然函数的每次调用，虽然数组 A 和 B 中的实际值在变化，但每个数组的维度保持不变。我想这可能有助于加快速度，因为相同的元素总是相乘。

Answer 1

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

Answer 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