如何高效计算两个Numpy数组的按行外积？

完全理解你的痛点——用Python循环+zip处理大数组确实会拖慢速度，Numpy本身就提供了向量化的解决方案，比你当前的实现高效得多，而且代码更简洁。

核心思路：利用Numpy广播机制 #

你的需求是对每个batch维度（b）的元素，计算A的行与B的行的外积。我们可以通过扩展维度让Numpy自动完成广播相乘，完全避免Python层面的循环：

import numpy as np
A = np.ones((10, 2))
B = np.ones((10, 6))

# 给A增加最后一个维度（形状变为(b, i, 1)），给B增加中间维度（形状变为(b, 1, o)）
R = A[:, :, None] * B[:, None, :]

assert R.shape == (10, 2, 6)

原理说明 #

• A[:, :, None]将A从(b, i)扩展为(b, i, 1)
• B[:, None, :]将B从(b, o)扩展为(b, 1, o)
• Numpy的广播机制会自动将这两个数组扩展为(b, i, o)的形状，对应位置相乘就等价于每个batch内的外积操作。

更直观的实现：使用np.einsum #

如果你觉得维度扩展的写法不够直观，可以用np.einsum（爱因斯坦求和约定），它能直接用字符串描述维度之间的关系：

R = np.einsum('bi,bo->bio', A, B)

语法解释 #

• 'bi,bo->bio'的含义是：
  1. 取A的b（batch）和i维度，B的b和o维度
  2. 输出一个包含b、i、o三个维度的数组
  3. 每个b对应的i和o维度的元素，就是A该行与B该行的外积

效率对比（以大数据集为例） #

我们用更大的数组测试三种方法的速度（借助IPython的%timeit）：

A_large = np.random.rand(10000, 100)
B_large = np.random.rand(10000, 200)

# 原方法
%timeit np.asarray([np.outer(a, b) for a, b in zip(A_large, B_large)])
# 广播方法
%timeit A_large[:, :, None] * B_large[:, None, :]
# einsum方法
%timeit np.einsum('bi,bo->bio', A_large, B_large)

测试结果通常会显示：广播和einsum方法的速度是原方法的几十到上百倍，因为它们是在C层面执行的向量化操作，完全避开了Python循环的开销。

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