如何在NumPy中实现矩阵数组与向量数组的无循环批量相乘？

当然可以用NumPy的内置操作实现，完全不需要显式循环，而且方法能无缝适配更高维度的批量数据。

基础场景：A.shape=(N,3,3)，V.shape=(N,3) #

有两种常用方案：

方案1：使用np.einsum #

通过维度下标直接定义运算关系，可读性极强：

import numpy as np

result = np.einsum('nij,ni->nj', A, V)

下标解释：

• nij对应A的维度（N个3×3矩阵）
• ni对应V的维度（N个3维向量）
• nj指定输出维度（N个3维结果向量）
  本质是对每个i，计算A[i]（3×3）与V[i]（3）的矩阵向量积，最后按N维度堆叠结果。

方案2：使用np.matmul（或@运算符） #

利用矩阵乘法的批量处理特性，先将V扩展为(N,3,1)的列向量，相乘后再压缩最后一维：

result = np.matmul(A, V[..., np.newaxis]).squeeze(-1)
# 或用@运算符简化
result = (A @ V[..., None]).squeeze(-1)

验证正确性

可以对比循环实现的结果：

N = 5
A = np.random.rand(N, 3, 3)
V = np.random.rand(N, 3)

# 循环实现
loop_result = np.array([A[i] @ V[i] for i in range(N)])

# einsum实现
einsum_result = np.einsum('nij,ni->nj', A, V)

# matmul实现
matmul_result = np.matmul(A, V[..., np.newaxis]).squeeze(-1)

print(np.allclose(loop_result, einsum_result))  # 输出True
print(np.allclose(loop_result, matmul_result))  # 输出True

高维场景：A.shape=(N,M,3,3)，V.shape=(N,M,3) #

上述两种方案完全适配，只需要调整维度下标或直接复用代码：

方案1：np.einsum调整下标 #

result = np.einsum('nmij,nmi->nmj', A, V)

下标nmij对应A的(N,M,3,3)维度，nmi对应V的(N,M,3)维度，输出nmj即(N,M,3)的结果数组。

方案2：np.matmul直接复用代码 #

因为matmul会自动识别前面的批量维度（N和M），无需修改核心逻辑：

result = np.matmul(A, V[..., np.newaxis]).squeeze(-1)
# 或
result = (A @ V[..., None]).squeeze(-1)

验证正确性

N = 2
M = 3
A = np.random.rand(N, M, 3, 3)
V = np.random.rand(N, M, 3)

# 循环实现
loop_result = np.array([[A[n,m] @ V[n,m] for m in range(M)] for n in range(N)])

# einsum实现
einsum_result = np.einsum('nmij,nmi->nmj', A, V)

# matmul实现
matmul_result = np.matmul(A, V[..., np.newaxis]).squeeze(-1)

print(np.allclose(loop_result, einsum_result))  # 输出True
print(np.allclose(loop_result, matmul_result))  # 输出True

方案对比 #

• np.einsum：优势是直观，通过下标就能清晰看到维度的对应关系，适合复杂的多维度运算场景；
• np.matmul/@：优势是性能更优，底层依赖优化过的BLAS/LAPACK实现，大规模数据下速度更快。

内容的提问来源于stack exchange，提问作者Alex V.