如何在NumPy中加速大量小规模协方差的高频计算？

1 替换低效的cov实现 #

原有np.cov会计算完整的2×2协方差矩阵，你仅需要取其中一个元素，存在大量冗余计算。直接通过协方差定义实现可以减少不必要的开销，性能提升至少3~5倍：

def cov(a, b):
    n = a.shape[0]
    return ((a - a.mean()) * (b - b.mean())).sum() / (n - 1)

如果你的场景使用有偏协方差，把分母的n-1改成n即可，和np.cov(a, b, bias=True)[0,1]结果一致。如果向量长度n是固定值，也可以直接写死分母，省掉每次取shape的开销。

2 预计算重复的逐元素乘积 #

你当前代码里大量重复计算了ix2*x2、ix2*y2、iy2*x2这类中间向量，提前统一计算一次，避免重复的逐元素运算开销：

def diff_cov_ridge(theta, tx, ty, img1, ix1, iy1, x1, y1, img2, ix2, iy2, x2, y2):
    ct = np.cos(theta)
    st = np.sin(theta)
    # 预计算所有重复用到的乘积向量
    ix2_x2 = ix2 * x2
    ix2_y2 = ix2 * y2
    iy2_x2 = iy2 * x2
    iy2_y2 = iy2 * y2
    ix1_x1 = ix1 * x1
    ix1_y1 = ix1 * y1
    iy1_x1 = iy1 * x1
    iy1_y1 = iy1 * y1
    tx_x1 = tx * x1
    tx_x2 = tx * x2
    tx_y1 = tx * y1
    tx_y2 = tx * y2
    ty_x2 = ty * x2
    ty_y2 = ty * y2

    eq1 = cov(img1, ix2_x2)
    eq2 = cov(img1, ix2_y2)
    eq3 = cov(img1, iy2_x2)
    eq4 = cov(img1, iy2_y2)
    eq5 = cov(img2, ix1_x1)
    eq6 = cov(img2, ix1_y1)
    eq7 = cov(img2, iy1_x1)
    eq8 = cov(img2, iy1_y1)
    eq9 = cov(ix2, ix1 * tx_x1)
    eq10 = cov(ix1, ix2 * tx_x2)
    eq11 = cov(ix1_y1, ix2 * tx)
    eq12 = cov(ix1, ix2 * tx_y2)
    eq13 = cov(ix1_x1, ix2_x2)
    eq14 = cov(ix1_x1, ix2_y2)
    eq15 = cov(ix1_y1, ix2_x2)
    eq16 = cov(ix1_y1, ix2_y2)
    eq17 = cov(ix1, iy2 * tx_x2)
    eq18 = cov(ix1, iy2 * tx_y2)
    eq19 = cov(ix1_x1, iy2 * ty)
    eq20 = cov(ix1_y1, iy2 * ty)
    eq21 = cov(ix1_x1, iy2_x2)
    eq22 = cov(ix1_x1, iy2_y2)
    eq23 = cov(ix1_y1, iy2_x2)
    eq24 = cov(ix1_y1, iy2_y2)
    eq25 = cov(ix2, iy1 * tx_x1)
    eq26 = cov(ix2, iy1 * tx_y1)
    eq27 = cov(iy1, ix2 * ty_x2)
    eq28 = cov(iy1, ix2 * ty_y2)
    eq29 = cov(ix2_x2, iy1_x1)
    eq30 = cov(ix2_y2, iy1_x1)
    eq31 = cov(ix2_x2, iy1_y1)
    eq32 = cov(ix2_y2, iy1_y1)
    eq33 = cov(iy1_x1, iy2 * ty)
    eq34 = cov(iy1, iy2 * ty_x2)
    eq35 = cov(iy1_y1, iy2 * ty)
    eq36 = cov(iy1, iy2 * ty_y2)
    eq37 = cov(iy1_x1, iy2_x2)
    eq38 = cov(iy1_x1, iy2_y2)
    eq39 = cov(iy1_y1, iy2_x2)
    eq40 = cov(iy1_y1, iy2_y2)
    # 后续逻辑保持不变

3 用Numba JIT编译获得数量级性能提升 #

如果你需要调用数百万次，上述纯NumPy版本依然有Python函数调用、逐次均值计算的开销，用Numba的@njit装饰器编译整个函数，性能可以再提升10~100倍，完全满足百万次调用的需求：

import numba
import numpy as np

@numba.njit(fastmath=True)
def cov_numba(a, b):
    n = a.shape[0]
    a_mean = a.mean()
    b_mean = b.mean()
    res = 0.0
    for i in range(n):
        res += (a[i] - a_mean) * (b[i] - b_mean)
    return res / (n - 1)

@numba.njit(fastmath=True)
def diff_cov_ridge_numba(theta, tx, ty, img1, ix1, iy1, x1, y1, img2, ix2, iy2, x2, y2):
    ct = np.cos(theta)
    st = np.sin(theta)
    # 所有中间预计算逻辑和上面版本一致，替换cov为cov_numba即可
    ix2_x2 = ix2 * x2
    ix2_y2 = ix2 * y2
    # 其余协方差计算逻辑同上

如果你的输入向量长度n是固定的小数值（比如小于16），还可以给Numba加上静态类型标注，编译后性能和手写C代码相当。

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