循环 np.einsum 多次...有更快的方法吗?
Looping over np.einsum many times... Is there a faster way?
我有一个似然函数,我正在尝试使用 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 中的实际值在变化,但每个数组的维度保持不变。 I'我想这可能有助于加快速度,因为 相同的元素总是相乘。
即使在小循环中使用 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
好吧,其中一个轴在 A
(第一个)和 B
(最后一个)中保持对齐,并且也在输出(最后一个)中保持对齐,并且是一个非常小的循环数4
。因此,我们可以简单地用 np.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
.
我有一个似然函数,我正在尝试使用 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 中的实际值在变化,但每个数组的维度保持不变。 I'我想这可能有助于加快速度,因为 相同的元素总是相乘。
即使在小循环中使用 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
好吧,其中一个轴在 A
(第一个)和 B
(最后一个)中保持对齐,并且也在输出(最后一个)中保持对齐,并且是一个非常小的循环数4
。因此,我们可以简单地用 np.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
.