Numpy 数组沿特定轴相乘

Question

我想沿特定轴将两个不同形状的 numpy 数组相乘，而不需要手动交换它们或在可能的情况下添加“虚拟”维度。

对于两个数组，假设形状为 (3,2,2) 的 A 和形状为 (3) 的 B，我想获得以下形式的乘积

np.swapaxes(np.swapaxes(A,0,2)*B),0,2)

或者替代地

A*B[:,None,None]

然而，对于较大的数组，与以 (2,2,3) 形状初始化 A 相比，这两种方法似乎都会显着增加过程的计算时间。但我不能这样做，因为我需要 A 的形状为 (3,2,2) 以便以后的乘法运算。

我可以避免双重交换轴或添加“无”维度来缩短计算时间吗？

Answer 1

在我的快速测试中，沿第三轴相乘实际上比沿最后一个轴相乘更快。

In [135]: A=np.ones((2,2,100,100,100)); B=np.ones((100))

In [136]: timeit A*B
17.9 ms ± 333 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [137]: timeit A*B[:,None,None]
16 ms ± 1.75 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)