Numpy实现多组矩阵克罗内克积求和的高效方法

问题描述 #

现有三组矩阵集合{A_i}、{B_i}、{C_i}，每组各包含n个矩阵：

• A_i维度为l×m，B_i维度为m×o，C_i维度为p×q
• 需要计算的表达式为：$\sum_{i=1}^n \text{kron}(A_i B_i, C_i)$（其中$A_i B_i$表示矩阵乘法，$\text{kron}$表示克罗内克积）

示例代码（原始实现） #

import numpy as np

A = np.arange(12).reshape(2,3,2)  # shape (n=2, l=3, m=2)
B = np.arange(12,24).reshape(2,2,3)  # shape (n=2, m=2, o=3)
C = np.arange(32).reshape(2,4,4)  # shape (n=2, p=4, q=4)

result = np.zeros((12,12))
for i in range(2):
    result += np.kron(A[i,:,:] @ B[i,:,:], C[i,:,:])

高效实现方案 #

原始实现通过Python循环逐个计算克罗内克积再累加，当n较大时会因Python循环的解释开销导致效率低下。我们可以利用numpy矢量化操作，结合张量广播和维度变换实现无循环的高效计算：

步骤说明 #

1. 批量计算矩阵乘法：用np.matmul直接对所有A_i和B_i做批量矩阵乘法，得到形状为(n, l, o)的张量AB。
2. 批量生成克罗内克积：通过维度扩展和广播，将AB与C的元素逐元素相乘，再调整维度得到每个kron(AB[i], C[i])的批量结果。
3. 批量求和：沿着矩阵集合的维度（第0维）求和，得到最终结果。

实现代码 #

import numpy as np

A = np.arange(12).reshape(2,3,2)
B = np.arange(12,24).reshape(2,2,3)
C = np.arange(32).reshape(2,4,4)

# 1. 批量计算所有A_i @ B_i
AB = np.matmul(A, B)  # shape (2, 3, 3)

# 2. 批量生成克罗内克积：利用维度扩展+广播+形状变换
AB_expanded = AB[:, :, np.newaxis, :, np.newaxis, np.newaxis]  # shape (2, 3, 1, 3, 1, 1)
C_expanded = C[:, np.newaxis, :, np.newaxis, :, np.newaxis]    # shape (2, 1, 4, 1, 4, 1)

# 广播相乘后调整为(n, l*p, o*q)的批量克罗内克积结果
kron_batch = (AB_expanded * C_expanded).reshape(A.shape[0], -1, AB.shape[2] * C.shape[2])

# 3. 沿第0维求和得到最终结果
result_fast = kron_batch.sum(axis=0)

# 验证与原始实现结果一致
assert np.allclose(result, result_fast)

效率优势 #

• 矢量化操作完全依托numpy底层的C优化（依赖BLAS/LAPACK库），避免了Python循环的解释开销，n越大，速度提升越显著。
• 批量操作的内存访问模式更连续，缓存命中率更高，进一步优化计算效率。

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