如何高效实现以样本量为第一维度的逐元素乘法

嘿，我来帮你搞定这个批量计算的效率问题！

你现在的情况是：有两个数组，A的形状是(100, 150, 1000)（100个样本，每个样本是150行1000列的矩阵），B的形状是(100, 150)（每个样本对应150个数值）。你已经能处理单个样本的(150,1000) * (150)，现在要高效批量处理这100个样本，其实用NumPy的广播机制就能完美解决，完全不用写Python循环，速度快到飞起！

具体实现步骤 #

• 把B从(100,150)扩展成(100,150,1)——可以用B[:, :, None]或者np.expand_dims(B, axis=2)，这样每个样本的150个数值就变成了150行1列的矩阵，刚好能和A中对应样本的150行1000列矩阵做逐元素乘法（每一行的1000个元素都乘上该行对应的B值）。
• 直接用*运算符相乘就行，NumPy会自动处理批量的样本，结果形状保持(100,150,1000)。

代码示例：

import numpy as np

# 假设A和B已经是NumPy数组
A = np.random.rand(100, 150, 1000)
B = np.random.rand(100, 150)

# 扩展B的维度，添加一个长度为1的第三维
B_expanded = B[:, :, None]  # 等价于np.expand_dims(B, axis=2)

# 批量逐元素乘法，结果形状为(100,150,1000)
result = A * B_expanded

为什么这个方法高效？ #

• 避免了手动写Python循环：Python循环在处理大数组时速度很慢，而NumPy的操作是底层用C实现的，能充分利用CPU的向量运算能力，速度提升非常明显。
• 广播机制不需要额外复制数据：扩展维度只是改变了数组的视图，不会占用额外的内存空间，内存效率也很高。

验证结果正确性 #

你可以拿第一个样本的结果和你之前手动计算的对比，确保没问题：

# 手动计算第一个样本的结果
manual_result = A[0] * B[0]

# 和批量计算的第一个样本对比
print(np.allclose(result[0], manual_result))  # 应该输出True

备注：内容来源于stack exchange，提问作者Lisa