如何用Numpy（非循环方式）实现矩阵的逐元素矩阵乘法？

核心思路 #

你的4维数组形状为 (M, N, 2, 2)，前两维是大矩阵的行列，后两维是每个位置的2x2方阵。要实现对应位置的矩阵乘法，本质是批量矩阵乘法——将前两维视为批量维度，对每个批量样本执行标准矩阵乘法，无需手动遍历元素。

NumPy 实现 #

直接使用np.matmul（或@运算符的批量版本），它会自动处理批量维度的广播，并行计算所有对应位置的矩阵乘法：

import numpy as np

# 示例输入
x = np.array([[  [[1,0], [0, 1]], [[2,2], [2, 1]]  ]])
y = np.array([[  [[1,3], [0, 1]], [[2,0], [0, 2]]  ]])

# 执行批量对应位置矩阵乘法
xy = np.matmul(x, y)

# 验证结果
print(np.array_equal(xy, np.array([[  [[1,3], [0, 1]], [[4,4], [4, 2]]  ]])))
# 输出: True

np.matmul遵循规则：若输入数组形状为(..., m, k)和(..., k, n)，结果形状为(..., m, n)。这里每个2x2方阵满足m=k=n=2，完全匹配需求。

CuPy 并行实现 #

CuPy兼容NumPy API，只需替换为cp即可利用GPU并行加速，底层自动调用优化的CUDA内核：

import cupy as cp

# 将数据转移到GPU
x_gpu = cp.array([[  [[1,0], [0, 1]], [[2,2], [2, 1]]  ]])
y_gpu = cp.array([[  [[1,3], [0, 1]], [[2,0], [0, 2]]  ]])

# 执行GPU并行批量矩阵乘法
xy_gpu = cp.matmul(x_gpu, y_gpu)

# 转回CPU（按需操作）
xy_cpu = cp.asnumpy(xy_gpu)

# 验证结果
print(np.array_equal(xy_cpu, np.array([[  [[1,3], [0, 1]], [[4,4], [4, 2]]  ]])))
# 输出: True

补充说明 #

• 若数组有更多批量维度（如(K, M, N, 2, 2)），matmul依然适用——所有前导维度会被视为批量维度，仅对最后两维执行矩阵乘法。
• 也可通过np.einsum('...ij,...jk->...ik', x, y)实现，但matmul通常有更优的底层性能优化。

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