Question

我有一个numpy数组 f ，长度为 n ，numpy矩阵 A ，大小为 n x < EM>米。我想打破 r 件中的 f 和 A f1 ，...， fr 和 A1 ，...， Ar ，然后进行计算 fi * Ai （数学意义上的向量x矩阵乘法）每个 fi 是一个行向量，其列数等于 Ai 的行数。结果将是行向量1 x m 。想法是连接所有这些行向量以形成矩阵 B = [[f1 * A1]，[f2 * A2]，...，[fr * Ar]] （注意这将是是一个大小为 r x m ）的矩阵。

假设已经定义了 f 和 A 。还假设片段的相应索引在列表 [0，d1，... dr] 中。例如， f1 = f [d [0]：d [1]] 和 f2 = f [d [1]：d [2]] ）。我正在使用以下代码来解决我的问题：

B = numpy.zeros([r,m])
for i in range(0,r):
    lower = d[i]
    upper = d[i+1]
    B[i,:] = f[lower:upper].dot(A[lower:upper,:])

问题是这段代码将在我的程序中多次计算。之前我听说Python for循环很慢，实际上我的代码的瓶颈就是这部分。我无法弄清楚如何对此进行矢量化，但我觉得这是可能的。我希望有人能指引我。感谢。

Answer 1

我认为这是一个有效的MCVE：

In [139]: f = np.arange(10)
In [140]: A = np.arange(20).reshape(10,2)
In [141]: f.dot(A)
Out[141]: array([570, 615])
In [142]: d = [0,2,5,10]
In [143]: for i,j in zip(d[:-1],d[1:]):
     ...:     print(f[i:j].dot(A[i:j,:]))
     ...:     
[2 3]
[58 67]
[510 545]

其中570 = 2+58+510。

In [145]: np.array([f[i:j].dot(A[i:j,:]) for i,j in zip(d[:-1],d[1:])])
Out[145]: 
array([[  2,   3],
       [ 58,  67],
       [510, 545]])

鉴于i:j切片的长度可能不同，可能很难“矢量化”。这是真正意义上的。我们可以隐藏迭代，但是以将所有迭代移动到编译代码的方式编写它将是棘手的。像cumsum这样的累积操作通常是最好的选择。我们经常不得不退后一步，从不同的角度看问题（而不是简单地删除循环）。

numba和cython通常用于加速迭代解决方案，但我不会进入这些解决方案。

如果d将数组分成相等的部分，我们可以使用重新整形来计算这些部分：

In [228]: A.shape
Out[228]: (10, 2)
In [229]: f.shape
Out[229]: (10,)
In [230]: f2 = f.reshape(2,5)
In [231]: A2 = A.reshape(2,5,2)

In [233]: np.einsum('ij,ijk->ik',f2,A2)
Out[233]: 
array([[ 60,  70],
       [510, 545]])

matmul运算符也可以运行，但它需要对维度进行一些调整：

In [236]: (f2[:,None,:]@A2)[:,0,:]
Out[236]: 
array([[ 60,  70],
       [510, 545]])

如果d将数组划分为几个大小，我想我们可以对常见大小进行分组，并对每个组执行上述重塑和einsum，但我还没有弄清楚细节：

In [238]: d = [0,2,5,7,10]
In [239]: np.array([f[i:j].dot(A[i:j,:]) for i,j in zip(d[:-1],d[1:])])
Out[239]: 
array([[  2,   3],
       [ 58,  67],
       [122, 133],
       [388, 412]])
In [240]: [f[i:j] for i,j in zip(d[:-1],d[1:])]
Out[240]: [array([0, 1]), array([2, 3, 4]), array([5, 6]), array([7, 8, 9])]

这里我们有两组，一组长度为2，另一组长度为3。

Answer 2

您可以使用np.add.reduceat：

# example data
>>> f = np.arange(10)
>>> A = np.arange(50).reshape(10, 5)
>>> split = [0, 3, 5, 10]
>>> 
# reduceat
>>> np.add.reduceat(f[:, None] * A, split[:-1], axis=0)
array([[  25,   28,   31,   34,   37],
       [ 125,  132,  139,  146,  153],
       [1275, 1310, 1345, 1380, 1415]])
>>> 
# double check against list comprehension
>>> [fi @ Ai for fi, Ai in zip(*map(np.split, (f, A), 2*(split[1:-1],)))]
[array([25, 28, 31, 34, 37]), array([125, 132, 139, 146, 153]), array([1275, 1310, 1345, 1380, 1415])]

如果由于blas加速矩阵乘法导致列表理解或@ hpaulj的解决方案或OP循环更快，我不会感到惊讶。

如何对此循环进行矢量化？

2 个答案: