Python下优化大矩阵L2范数计算：如何避免嵌套循环？

你遇到的问题非常典型——Python嵌套循环在处理超大维度矩阵时，因为解释器的逐次调用开销，会导致速度慢到无法接受。好在L2范数的计算可以通过数学推导转化为向量化矩阵运算，完全不需要循环，而且能利用底层线性代数库（比如BLAS）的并行优化，速度提升几个数量级。

核心思路：把L2范数拆解为矩阵运算 #

先回忆L2范数的定义：对于A的第i列向量a_i（1024维）和B的第j列向量b_j（1024维），C[i][j] = ||a_i - b_j||₂。我们可以把平方后的范数展开：

||a_i - b_j||₂² = ||a_i||₂² + ||b_j||₂² - 2a_iᵀb_j

这样整个矩阵C的计算就可以拆成几步纯矩阵操作，完全避开循环。

具体实现（以NumPy为例） #

import numpy as np

# 假设A是(1024, 307200)，B是(1024, 50)
A = np.random.randn(1024, 307200)
B = np.random.randn(1024, 50)

# 1. 计算A各列的L2范数平方 (形状: (1, 307200))
norm_A_sq = np.sum(A ** 2, axis=0, keepdims=True)
# 2. 计算B各列的L2范数平方 (形状: (1, 50))
norm_B_sq = np.sum(B ** 2, axis=0, keepdims=True)
# 3. 计算A转置与B的矩阵乘积 (形状: (307200, 50))
A_T_B = A.T @ B
# 4. 利用广播自动扩展维度，计算平方差
sq_diff = norm_A_sq.T + norm_B_sq - 2 * A_T_B
# 5. 开根号得到L2范数，用max避免浮点误差导致的负数
C = np.sqrt(np.maximum(sq_diff, 0.0))

为什么这比嵌套循环快？ #

• 所有运算都是向量化操作，底层由C实现的BLAS库执行，会自动利用CPU的多核并行甚至SIMD指令，效率远高于Python循环。
• 彻底避免了Python解释器在循环中反复调用函数的开销——这是嵌套循环慢的核心原因。

其他框架的适配（比如PyTorch） #

如果你用深度学习框架处理张量，思路完全一致，还能轻松切换到GPU加速：

import torch

A = torch.randn(1024, 307200)
B = torch.randn(1024, 50)

norm_A_sq = torch.sum(A ** 2, dim=0, keepdim=True)
norm_B_sq = torch.sum(B ** 2, dim=0, keepdim=True)
A_T_B = A.T @ B
sq_diff = norm_A_sq.T + norm_B_sq - 2 * A_T_B
C = torch.sqrt(torch.clamp(sq_diff, min=0.0))

# 移到GPU加速只需一行
A, B = A.cuda(), B.cuda()

注意事项 #

• 数值稳定性：由于浮点运算的精度问题，sq_diff可能会出现极小的负数，所以一定要用np.maximum或torch.clamp把数值限制在非负范围后再开根号。
• 内存占用：A.T @ B的形状是(307200, 50)，按float32计算仅占用约60MB，完全在常规内存范围内，不用担心溢出。

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