如何用Numpy高效实现3D张量与2D张量的自定义乘法？

首先明确你的运算本质：对于形状为 (N, M, K) 的3D张量，和形状为 (M, K) 的2D张量，计算每个第一维度切片（(M,K)）与2D张量的逐行点积，最终得到形状为 (N, M) 的2D张量。

方法1：使用 np.einsum（推荐，简洁且自动调用BLAS优化） #

np.einsum 可以直接用爱因斯坦求和符号描述运算逻辑，Numpy会自动将其优化为底层BLAS操作（如GEMV或GEMM），实现高效计算：

import numpy as np

# 示例数据
N, M, K = 100, 3, 3
3D_tensor = np.random.rand(N, M, K)
2D_tensor = np.array([[10, 20, 30], [40, 50, 60], [70, 80, 90]])

# 核心运算
result = np.einsum('nmk,mk->nm', 3D_tensor, 2D_tensor)

解释：

• 'nmk,mk->nm' 表示：对每个 n（3D张量的第一维度）、m（行维度），将 3D_tensor[n,m,k] 与 2D_tensor[m,k] 相乘后对 k（列维度）求和，最终输出形状为 (N,M) 的张量。

方法2：广播 + 求和（同样利用BLAS加速） #

通过广播将2D张量扩展为与3D张量匹配的维度，再做元素级相乘后求和，Numpy会自动优化该过程：

# 核心运算：将2D张量扩展为(1,M,K)，与(N,M,K)广播相乘后沿最后一维求和
result = np.sum(3D_tensor * 2D_tensor[None, :, :], axis=-1)

验证运算正确性 #

用你给出的示例手动计算验证：

# 示例切片
slice_0 = np.array([[1,2,3], [4,5,6], [7,8,9]])
# 计算单行点积
row0 = slice_0[0] @ 2D_tensor[0]  # 1*10 + 2*20 +3*30 = 140
row1 = slice_0[1] @ 2D_tensor[1]  #4*40 +5*50 +6*60= 730
row2 = slice_0[2] @ 2D_tensor[2]  #7*70 +8*80 +9*90= 1900
# 用einsum计算
test_result = np.einsum('mk,mk->m', slice_0, 2D_tensor)
print(test_result)  # 输出 [140 730 1900]，与手动计算一致

性能说明 #

这两种方法都会调用Numpy底层的BLAS优化实现（如OpenBLAS、MKL等），避免了Python层面的循环，运算效率远高于手动遍历切片。对于大规模张量，np.einsum 通常会生成更优的BLAS调用路径，推荐优先使用。

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