python - 如何使用广播来加速这段代码？

Question

我有以下多维数组。第一个轴表示一个 3 维向量。我想为每个计算 3×3 矩阵 x⋅x'。

我目前的解决方案：

arr.shape
# (3, 64, 64, 33, 187)

dm = arr.reshape(3,-1)
dm.shape
# (3, 25276416)

cov = np.empty((3,3,dm.shape[1]))
cov.shape
# (3, 3, 25276416)

这个 for 循环遍历所有 25,276,416 个元素，大约需要 1 或 2 分钟。

for i in range(dm.shape[1]):
    cov[...,i] = dm[:,i].reshape(3,1).dot(dm[:,i].reshape(1,3))

cov = cov.reshape((3,) + arr.shape)
cov.shape
# (3, 3, 64, 64, 33, 187)

Answer 1

那么你并没有真正减少使用矩阵乘法的任何轴，np.dot它只是在那里广播元素乘法。所以，你可以简单地使用NumPy broadcasting整个事情，就像这样 -

cov = dm[:,None]*dm

或者直接使用它arr来避免创建dm和所有重塑，就像这样 -

cov = arr[:,None]*arr