如何无循环实现3D矩阵与2D矩阵的批量乘法运算？

当然有办法摆脱循环！NumPy提供了多种向量化操作的方式，既能让代码更简洁，又能利用底层优化大幅提升运算效率。下面给你介绍几种常用的方法：

1. 使用np.einsum（最直观，对应数学逻辑） #

einsum通过爱因斯坦求和约定精确描述运算逻辑，可读性拉满：

import numpy as np

A = np.ones((6,6,10))
B = np.ones((6,10))
mat = np.einsum('ijk,jk->ik', A, B)

这里的ijk,jk->ik可以这么理解：

• ijk对应A的三个维度：第0轴（6行）、第1轴（6列）、第2轴（10个批次）
• jk对应B的两个维度：第0轴（6行）、第1轴（10个批次）
• ->ik表示保留A的第0轴（i）和两个数组的批次轴（k），对中间的j轴（维度为6）求和——完美对应每个批次里6×6矩阵和6维向量的乘法运算。

2. 调整维度+广播矩阵乘法 #

给B增加一个维度，让它和A的维度匹配，借助NumPy的广播特性完成批量运算：

mat = (A @ B[..., None]).squeeze(-1)

• B[..., None]把B的形状从(6,10)变成(6,1,10)，每个批次的向量就变成了6×1的列向量
• A（(6,6,10)）和调整后的B做矩阵乘法，每个批次得到6×1的结果，整体形状是(6,1,10)
• squeeze(-1)去掉最后一个多余维度，最终得到(6,10)的目标结果。

3. 使用np.tensordot（张量点积） #

tensordot可以指定两个张量的哪些维度做点积，非常适配这种批量矩阵-向量乘法场景：

mat = np.tensordot(A, B, axes=([1], [0]))

这里axes=([1], [0])表示：对A的第1轴（6列）和B的第0轴（6行）做点积，剩下的A的第0轴（6行）和两个数组的批次轴组合起来，刚好得到(6,10)的结果。

验证结果一致性 #

你可以用随机数据测试这几种方法和原始循环的结果是否完全一致：

# 生成随机测试数据
A = np.random.rand(6,6,10)
B = np.random.rand(6,10)

# 原始循环实现
mat_loop = np.zeros((6,10))
for i in range(10):
    mat_loop[:,i] = A[:,:,i]@B[:,i]

# 验证三种方法
print(np.allclose(mat_loop, np.einsum('ijk,jk->ik', A, B)))  # 输出True
print(np.allclose(mat_loop, (A @ B[..., None]).squeeze(-1)))  # 输出True
print(np.allclose(mat_loop, np.tensordot(A, B, axes=([1], [0]))))  # 输出True

这些向量化方法在数据量越大时优势越明显——NumPy底层是C优化实现，比Python循环快得多。

内容的提问来源于stack exchange，提问作者arundathi-c