Numpy下3D数组与2D数组逐元素乘求和的无循环高效实现

实现方案 #

你可以通过numpy的广播机制、张量运算等向量化操作完全替代Python级for循环，运算效率会有明显提升，以下是几种等价实现：

方法1：广播后沿空间维度求和（最直观易读）

利用numpy的广播规则给二维kernel扩充通道维度，自动匹配3通道的输入张量，再对高、宽两个空间维度求和即可：

# 扩充kernel为 (5,5,1) 后自动广播到3个通道逐元素相乘，对前两维求和
result = (kernel[..., None] * portion).sum(axis=(0, 1))

其中kernel[..., None]等价于kernel[:, :, np.newaxis]，作用是给kernel新增最后一个维度，适配输入的通道维度。

方法2：使用tensordot指定收缩维度（运算效率最优）

直接通过张量点积运算，指定要做乘加收缩的维度，无需手动广播：

# 将kernel的(0,1)维度和portion的(0,1)维度做乘加收缩
result = np.tensordot(kernel, portion, axes=((0, 1), (0, 1)))

方法3：使用einsum实现（灵活度最高）

通过爱因斯坦求和约定直接定义运算规则，代码更简洁：

# 定义运算规则：ij维度的kernel 和 ijk维度的portion 乘加后输出k维度的结果
result = np.einsum('ij,ijk->k', kernel, portion)

结果一致性验证 #

你可以通过以下代码验证以上实现和你原有循环的输出完全一致：

import numpy as np

# 生成测试数据
portion = np.random.rand(5, 5, 3)
kernel = np.random.rand(5, 5)

# 原循环实现
res_loop = np.array([(kernel * portion[:, :, c]).sum() for c in range(3)])
# 广播实现
res_broad = (kernel[..., None] * portion).sum(axis=(0, 1))
# tensordot实现
res_tensordot = np.tensordot(kernel, portion, axes=((0,1), (0,1)))
# einsum实现
res_einsum = np.einsum('ij,ijk->k', kernel, portion)

print(np.allclose(res_loop, res_broad))      # 输出 True
print(np.allclose(res_loop, res_tensordot))  # 输出 True
print(np.allclose(res_loop, res_einsum))     # 输出 True

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