如何用numpy.einsum或einops简化张量与矩阵的批量乘法运算？

问题说明 #

给定两个输入：

• T：维度为 (H x W x C) 的张量，必要时可转换为 (H x W x C x 1)
• M：维度为 (C x C) 的矩阵

需要计算一个维度为 (H x W x C) 的结果张量，其中每个 (1 x 1 x C) 的切片与矩阵M执行矩阵乘积。现有嵌套循环实现如下：

result = np.zeros_like(T)
for row in range(H):
    for col in range(W):
        result[row, col, ...] = M @ T[row, col, ...]

以下是三种更高效的实现方式：

1. Numpy原生向量化实现 #

利用Numpy的广播机制和矩阵乘法运算符，避免显式循环：

# 方法1：扩展T的维度为(H,W,C,1)，执行矩阵乘法后压缩最后一维
result = (M @ T[..., np.newaxis]).squeeze(-1)

# 方法2：转置维度后批量计算，再还原维度
result = (M @ T.transpose(2, 0, 1)).transpose(1, 2, 0)

Numpy的向量化运算基于底层C实现，批量处理效率远高于Python层的嵌套循环。

2. 使用numpy.einsum #

通过爱因斯坦求和符号直接描述维度对应关系，可读性与效率兼顾：

# 维度映射：M的(C_in,C_out) 对应 T的(H,W,C_in)，输出(H,W,C_out)
result = np.einsum('ij,hwj->hwi', M, T)

einsum会自动优化运算路径，在多数场景下性能与原生向量化实现持平，且维度关系更直观。

3. 使用einops #

einops提供更贴近自然语言的维度操作语法，降低维度变换的出错概率：

from einops import einsum, rearrange

# 方法1：直接用einops的einsum描述运算
result = einsum(T, M, 'h w c_in, c_in c_out -> h w c_out')

# 方法2：先重排维度执行矩阵乘法，再还原维度
result = rearrange(M @ rearrange(T, 'h w c -> c h w'), 'c h w -> h w c')

einops底层依赖Numpy的高效实现，性能与原生方法一致，但代码可读性更强。

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