优化含点积的Python for循环代码执行时间

你的代码存在两个核心问题导致执行效率低下：一是Python层面的for循环遍历每个样本，无法利用NumPy的向量化运算优势；二是使用了已弃用的np.matrix类（官方推荐使用np.ndarray）。此外代码还存在几处语法错误：Hy[i].transpose需改为Hy[i].transpose()，多处括号未闭合。

优化思路 #

将循环内的逐样本运算转换为批量矩阵运算，利用NumPy的广播机制和底层C实现来加速计算。原代码的数学逻辑是累加每个样本的 ( H_i^T (y_i - H_i\mu_0)(y_i - H_i\mu_0)^T H_i )，再添加正则项，我们可以通过维度重塑和批量矩阵操作直接完成这个累加。

优化后的代码 #

import numpy as np

# 先确保输入数组的形状符合批量运算要求：
# y: 形状为 (n, m) 或 (n, m, 1)，n是样本数，m是特征维度
# Hy: 形状为 (n, m, k)，每个样本对应一个m×k的矩阵
# Mu0: 形状为 (k,) 或 (k, 1)

# 统一调整维度，适配批量矩阵运算
y = y.reshape(n, -1, 1)       # 转为 (n, m, 1)
Hy = Hy.reshape(n, -1, k)
Mu0 = Mu0.reshape(-1, 1)      # 转为 (k, 1)

# 批量计算所有样本的残差
residuals = y - Hy @ Mu0      # 形状：(n, m, 1)

# 计算残差的外积（每个样本对应一个m×m矩阵）
residual_outer = residuals @ residuals.transpose(0, 2, 1)  # 形状：(n, m, m)

# 批量计算每个样本的贡献矩阵，然后沿样本维度求和
S = np.sum(Hy.transpose(0, 2, 1) @ residual_outer @ Hy, axis=0)  # 形状：(k, k)

# 添加L2正则项
S += np.eye(k) / 10**6

优化效果说明 #

• 移除Python循环：所有运算都在NumPy的底层C代码中执行，彻底消除了Python循环的性能开销，当样本数n较大时，速度提升会非常明显。
• 使用ndarray替代np.matrix：ndarray更灵活，支持任意维度，且避免了np.matrix的隐式转置等容易混淆的行为。
• 批量运算：通过维度重塑将所有样本的计算合并为一次矩阵操作，充分利用CPU的缓存和向量运算能力。

如果你的数据集非常大（n极大导致内存不足），可以考虑分块处理样本，每次计算一部分样本的贡献并累加，但在大多数场景下，上述向量化版本已经足够高效。

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