Question

我有一个要用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中的实际值都在变化，而每个数组的维数保持固定。可能会有助于加快速度，因为相同的元素总是被相乘在一起。

Answer 1

即使在小循环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

Answer 2

其中一个轴在A（第一个）和B（最后一个）中保持对齐，并且在输出中也保持（最后一个）对齐，并且是{{ 1}}。因此，我们可以简单地用np.tensordot遍历那个，以减少张量和。使用如此大的数据集时，4减少内存拥塞的好处可能会克服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)的最后一个轴求和，然后我们会看到{ {1}}版-

因此，在这种情况下，最好使用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的数组初始化，并在每次调用log-拟议的似然函数，其循环使用B并更新输出out = np.empty(A.shape[1:3] + B.shape[1:])。

多次遍历np.einsum ...有没有更快的方法？

2 个答案:

特定于给定的问题