Python中3D数组逐元素乘法实现与性能优化问题

优化后实现代码 #

import numpy as np

# 原始输入参数（l=2, m=b=4 示例）
g_l3 = np.array([[1, 4, 5],[2, 6, 7]])
A_lm = np.arange(1, 9, 1).reshape(2, 4)
B_lb = np.arange(7, 15, 1).reshape(2, 4)

# 步骤1：用einsum一步完成张量乘积+维度求和，无冗余中间数组
D_mb33 = np.einsum('li, lj, lm, lb -> mbij', g_l3, g_l3, A_lm, B_lb, optimize='greedy')

# 步骤2：用transpose+reshape替代两次concatenate，零拷贝完成维度重排
D = D_mb33.transpose(0, 2, 1, 3).reshape(3 * A_lm.shape[1], 3 * B_lb.shape[1])

核心优化逻辑 #

• 替换广播乘加为np.einsum运算：原实现中g_l33、AB_lmb以及后续的广播相乘会生成大量冗余中间数组，当l、m规模达到1e4级别时，中间数组内存占用会飙升到数十GB，触发内存分页后性能暴跌。np.einsum会直接按照张量缩并规则完成运算，无需存储全量中间结果，同时会自动调用BLAS底层优化指令，单核性能通常是手动广播的3~10倍，大尺寸下优势更明显。
• 替换拼接操作为维度重排+重塑：原实现的两次np.concatenate需要对全量数据进行拷贝，当m达到1e4时，输出数组尺寸为3e4×3e4共9e8个元素，拼接开销极大。transpose调整维度顺序后直接reshape属于视图操作，不会复制数组数据，几乎没有额外开销。

额外性能提升建议 #

• 如果使用的numpy版本带MKL或OpenBLAS加速库，optimize='greedy'参数会自动选择最优张量缩并顺序，进一步降低运算复杂度。
• 若硬件支持GPU，可替换为cupy.einsum实现GPU加速，针对1e4级别的m/b参数，运算速度可提升1~2个数量级。

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