Question

我有以下维度的这些变量：

A   - (3,)
B   - (4,)
X_r - (3,K,N,nS)
X_u - (4,K,N,nS)
k - (K,)

我希望为每个可能的(A.dot(X_r[:,:,n,s])*B.dot(X_u[:,:,n,s])).dot(k)和n计算s，我现在的做法如下：

np.array([[(A.dot(X_r[:,:,n,s])*B.dot(X_u[:,:,n,s])).dot(k) for n in xrange(N)] for s in xrange(nS)]) #nSxN

但这是超级慢，我想知道是否有更好的方法，但我不确定。

然而，我正在做另一项计算，我相信它可以进行优化：

np.sum(np.array([(X_r[:,:,n,s]*B.dot(X_u[:,:,n,s])).dot(k) for n in xrange(N)]),axis=0)

在这个中我创建了一个numpy数组，只是为了在一个轴上求和，然后丢弃数组。如果这是1-D中的列表，我会使用reduce并对其进行优化，我应该为numpy数组使用什么？

Answer 1

使用少量np.einsum来电 -

# Calculation of A.dot(X_r[:,:,n,s])
p1 = np.einsum('i,ijkl->jkl',A,X_r)

# Calculation of B.dot(X_u[:,:,n,s])
p2 = np.einsum('i,ijkl->jkl',B,X_u)

# Include .dot(k) part to get the final output
out = np.einsum('ijk,i->kj',p1*p2,k)

关于第二个例子，这解决了它：

p1 = np.einsum('i,ijkl->jkl',B,X_u)#OUT_DIM - (k,N,nS)
sol = np.einsum('ijkl,j->il',X_r*p1[None,:,:,:],k)#OUT_DIM (3,nS)

Answer 2

您可以使用dot来表示更高维度的矩阵乘法，但运行索引必须是最后两个。当我们重新排序你的矩阵时

X_r_t = X_r.transpose(2,3,0,1)
X_u_t = X_u.transpose(2,3,0,1)

我们获取您的第一个表达

res1_imp = (A.dot(X_r_t)*B.dot(X_u_t)).dot(k).T # shape nS x N

和第二个表达

res2_imp = np.sum((X_r_t * B.dot(X_u_t)[:,:,None,:]).dot(k),axis=0)[-1]

<强>计时

Divakars解决方案在我的计算机上10000 loops, best of 3: 21.7 µs per loop

我的解决方案提供10000 loops, best of 3: 101 µs per loop

修改

我的上部计时包括两个表达式的计算。当我只包含第一个表达式（作为Divakar）时，我获得了10000 loops, best of 3: 41 µs per loop ...它仍然较慢但更接近他的时间

使用Numpy的最快方式 - 多维总和和产品

2 个答案: