Python百万次循环中小数组沿最后轴计算点积的高效方法？

1. 替换einsum为更轻量的逐元素乘+求和 #

einsum的字符串解析和通用调度逻辑在处理小数组时，额外开销占比极高。直接用逐元素相乘后沿最后轴求和，能跳过这些冗余开销，速度提升明显：

# 假设每次循环的两个输入是a（类数组，形状(N, K)）和b（类数组，形状(N, K)）
result = (a * b).sum(axis=-1)
# 等价写法：
result = np.sum(a * b, axis=-1)

针对你的测试场景，先把列表转成numpy数组（避免循环中遍历列表的额外开销）：

test_arr = np.arange(200).reshape(20, 10)
test_list_arr = np.array(test_list).T  # 转成(20,10)，与test_arr形状匹配
result = (test_arr * test_list_arr).sum(axis=1)

2. 向量化整个循环（最优方案） #

如果百万次的所有输入可以提前整理成高维数组，直接消除Python循环的开销：
比如把百万次的test_arr类对象堆叠成all_arrays = np.stack([arr1, arr2, ..., arr1e6])（形状(1e6, 20, 10)），对应的all_lists也堆叠成相同形状，然后一次性计算：

all_results = (all_arrays * all_lists).sum(axis=-1)

numpy会在C层完成所有循环，速度比Python循环快几十到上百倍。

3. 用Numba JIT加速循环 #

如果无法提前整理输入（比如每次循环的输入是动态生成的），用Numba编译循环内的逻辑，消除Python解释器的开销：

import numba

@numba.njit  # 编译成机器码，首次调用后无解释器开销
def compute_dot(a, b):
    # a和b是(N, K)的数组，返回(N,)的点积结果
    n = a.shape[0]
    res = np.empty(n, dtype=a.dtype)
    for i in range(n):
        total = 0.0
        for j in range(a.shape[1]):
            total += a[i, j] * b[i, j]
        res[i] = total
    return res

# 百万次循环调用
for _ in range(1_000_000):
    result = compute_dot(test_arr, test_list_arr)

Numba会把函数编译成高效的机器码，循环内的操作没有Python的额外开销，小数组计算速度远超纯numpy方案。

性能对比参考 #

在20x10的小数组单次计算中：

• einsum单次调用开销约几十微秒
• (a*b).sum(axis=-1)单次调用开销约几微秒
• Numba编译后的函数单次调用开销仅几百纳秒

百万次循环下，性能差距会被放大到数秒甚至数十秒的级别。

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