Numpy中批量执行大量矩阵运算的最快向量化实现方式是什么？

核心思路是将两组小矩阵分别堆叠为三维数组，利用Numpy内置的广播矩阵乘法能力，所有运算在底层C/BLAS层面执行，完全规避Python层循环开销，是纯Numpy环境下的最快实现。

实现代码 #

import numpy as np

# 配置参数
N = 100000  # 矩阵对数量
n = 2       # 单个小矩阵的尺寸

# 模拟输入：如果你原始数据是Python列表形式，用np.stack转为三维数组
a_list = [np.random.rand(n, n) for _ in range(N)]
b_list = [np.random.rand(n, n) for _ in range(N)]
a_stack = np.stack(a_list)  # 形状为 (N, n, n)
b_stack = np.stack(b_list)  # 形状为 (N, n, n)

# 逐对执行矩阵乘法，结果形状为 (N, n, n)，第i个输出即为 a_i @ b_i
result = a_stack @ b_stack

如果你可以直接生成三维格式的输入数据，不需要经过Python列表中转，可以省略np.stack步骤，进一步提升效率。

性能优势 #

• 运算效率相比原生Python循环、np.vectorize方案提升100倍以上，底层自动调用优化BLAS库利用CPU SIMD指令加速
• 内存占用可控，仅需要存储3个(N, n, n)数组，n=2时10万组数据总内存占用不到10MB，完全不会出现块对角矩阵方案的内存爆炸问题

可选等价写法 #

你也可以用np.matmul或者np.einsum实现等价效果，性能差异极小：

# 写法1：np.matmul
result = np.matmul(a_stack, b_stack)

# 写法2：np.einsum
result = np.einsum('ijk,ikl->ijl', a_stack, b_stack)

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