如何实现3D矩阵与2D矩阵的对应行Numpy点积运算？

我有一个形状为(2, 10, 3)的3D数组normal，以及一个形状为(2, 3)的2D数组t，具体数据如下：

print(t) #2d array
# Output:
[[1.003   2.32    3.11   ]
 [1.214   5.32    2.13241]]

print(normal) #3d array
# Output:
[[[0.69908573 0.0826756  0.84485978]
  [0.51058213 0.4052637  0.5068118 ]
  [0.45974276 0.25819549 0.10780089]
  [0.27484999 0.33367648 0.128262  ]
  [0.35963389 0.77600065 0.89393939]
  [0.46937506 0.59291623 0.06620307]
  [0.87603987 0.44414505 0.83394174]
  [0.83186093 0.62491876 0.38160734]
  [0.96819897 0.80183442 0.75102768]
  [0.54182908 0.19403844 0.07925769]]

 [[2.82248573 3.2341756  0.96825978]
  [2.63398213 3.5567637  0.6302118 ]
  [2.58314276 3.40969549 0.23120089]
  [2.39824999 3.48517648 0.251662  ]
  [2.48303389 3.92750065 1.01733939]
  [2.59277506 3.74441623 0.18960307]
  [2.99943987 3.59564505 0.95734174]
  [2.95526093 3.77641876 0.50500734]
  [3.09159897 3.95333442 0.87442768]
  [2.66522908 3.34553844 0.20265769]]]

需要让2D数组t中的每一行，与3D数组normal中对应的矩阵进行点积运算，最终得到一个形状为(2, 10)的数组——第n行包含t的第n行与normal的第n个矩阵中10个向量的点积结果。预期输出如下：

[0.62096458 0.62618459 0.37528887 0.5728386  1.19634398 0.79620507
 1.997884   0.75229492 1.2236496  0.4210626 ]
[2.96347746 3.30738892 3.50596579 4.93082295 5.33811805 4.44872493
 7.33480393 4.19173472 4.7406248  7.83229689]

解决方案 #

以下是三种基于NumPy的实现方式，都能满足需求：

方法1：使用np.einsum #

通过爱因斯坦求和符号清晰指定维度运算关系，适合多维数组的点积场景：

import numpy as np

# 构造示例数组
t = np.array([[1.003, 2.32, 3.11], [1.214, 5.32, 2.13241]])
normal = np.array([
    [[0.69908573, 0.0826756, 0.84485978],
     [0.51058213, 0.4052637, 0.5068118],
     [0.45974276, 0.25819549, 0.10780089],
     [0.27484999, 0.33367648, 0.128262],
     [0.35963389, 0.77600065, 0.89393939],
     [0.46937506, 0.59291623, 0.06620307],
     [0.87603987, 0.44414505, 0.83394174],
     [0.83186093, 0.62491876, 0.38160734],
     [0.96819897, 0.80183442, 0.75102768],
     [0.54182908, 0.19403844, 0.07925769]],
    [[2.82248573, 3.2341756, 0.96825978],
     [2.63398213, 3.5567637, 0.6302118],
     [2.58314276, 3.40969549, 0.23120089],
     [2.39824999, 3.48517648, 0.251662],
     [2.48303389, 3.92750065, 1.01733939],
     [2.59277506, 3.74441623, 0.18960307],
     [2.99943987, 3.59564505, 0.95734174],
     [2.95526093, 3.77641876, 0.50500734],
     [3.09159897, 3.95333442, 0.87442768],
     [2.66522908, 3.34553844, 0.20265769]]
])

result = np.einsum('ij,ijk->ik', t, normal)
print(result)

'ij,ijk->ik'表示：对t的第i行第j列元素，与normal的第i个矩阵第k个向量的第j列元素相乘后求和，最终得到第i行第k列的结果，完全匹配需求。

方法2：使用np.matmul #

调整t的维度后进行矩阵乘法，再去除多余维度：

# 将t从(2,3)转为(2,1,3)，与转置后的normal(2,3,10)做矩阵乘法，最后压缩维度
result = np.matmul(t[:, np.newaxis, :], normal.transpose(0,2,1)).squeeze()
print(result)

更简洁的写法：

result = np.matmul(t[:, None], normal.swapaxes(1,2)).squeeze()

方法3：广播相乘后求和 #

利用NumPy广播机制，让对应元素相乘后对最后一个维度求和：

# 扩展t的维度为(2,1,3)，与normal(2,10,3)广播相乘，再对第3维求和
result = (t[:, np.newaxis, :] * normal).sum(axis=-1)
print(result)

这种方式逻辑直观，容易理解多维数组的运算过程。

以上三种方法都能输出符合预期的(2,10)形状数组。

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