如何快速向量化3D NumPy数组的乘法运算？

问题描述 #

给定三个NumPy数组F、B、Q，形状如下：

F.shape = (d,l)
B.shape = (l,d)
Q.shape = (l,d,d)

其中l远大于d，需要实现如下循环运算的向量化加速：

A = np.zeros(shape=(d,d))
for ii in range(0,d):
    for jj in range(0,d):
        A[ii,jj] = np.sum(F[ii,:]*B[:,jj]*Q[:,ii,jj])

已知当Q为二维方阵时，可通过A = np.multiply(np.matmul(F,B),Q)实现，求针对三维Q的单行向量化写法。

向量化实现方案 #

最优单行写法（einsum） #

利用NumPy的einsum函数直接描述维度运算关系，完全匹配原始循环逻辑，写法简洁高效：

A = np.einsum('il,lj,lij->ij', F, B, Q)

其他等价实现 #

• 广播结合求和的写法：

A = (F[:, None, :] * B.T[None, :, :] * Q.transpose(1, 2, 0)).sum(axis=-1)

• tensordot实现：

A = np.tensordot(F[:, None, :] * B, Q, axes=([2], [0]))

方案说明 #

原始循环中，每个A[ii,jj]是三个长度为l的一维数组（F[ii,:]、B[:,jj]、Q[:,ii,jj]）逐元素相乘后求和。einsum通过下标直接定义了维度间的运算规则：

• il对应F的(d,l)维度
• lj对应B的(l,d)维度
• lij对应Q的(l,d,d)维度
• ->ij指定最终输出(d,d)的数组，沿l维度完成求和操作

这种写法充分利用NumPy的底层优化，在l远大于d的场景下，相比显式循环能获得数量级的性能提升。

内容的提问来源于stack exchange，提问作者stunted_inference123212321