如何避免轴交换或虚拟维度提升Numpy数组指定轴乘法效率?
如何在不手动交换轴或添加虚拟维度的前提下提升Numpy数组指定轴乘法的效率?
我希望在不手动交换轴或添加虚拟维度的前提下,对不同形状的Numpy数组沿指定轴进行乘法运算。
以形状为(3,2,2)的数组A和形状为(3)的数组B为例,我需要得到类似np.swapaxes(np.swapaxes(A,0,2)*B),0,2)或A*B[:,None,None]的运算结果,但对于更大的数组,这两种方法都会显著增加计算时间——尽管将A初始化为(2,2,3)形状能提升效率,但后续乘法操作需要A保持(3,2,2)形状,因此无法采用该方式。
能否通过避免双重轴交换或添加虚拟维度来提升计算效率?
编辑补充
感谢各位的解答,以下是我实际计算的详细示例:
import numpy as np from time import time A = np.random.rand(100,100,2,20,100) B = np.random.rand(100) C = np.random.rand(2,20,100) D = np.random.rand(100,100,2,20,100) start = time() for k in range(98,-1,-1): A[k,:,:,:,:] = np.pad(C[:,:,k:],((0,0),(0,0),(0,k)), constant_values=1)*np.pad(B[k:],((0,k)), constant_values=1)*A[k+1,:,:,:,:] + D[k,:,:,:,:] end = time() print('straight forward product: ', end-start) start = time() for k in range(98,-1,-1): A[k,:,:,:,:] = np.pad(B[k:],((0,k)), constant_values=1)[:,None,None,None]*A[k,:,:,:,:]*np.pad(C[:,:,k:],((0,0),(0,0),(0,k)), constant_values=1)*A[k+1,:,:,:,:] + D[k,:,:,:,:] end = time() print('None Dummy-Dimensions: ', end-start)
对我而言第二种方法耗时几乎翻倍,但我现在认为这实际上是乘法运算符的顺序导致的,而非轴的乘法对应方式。
再编辑
我知道这两种计算的结果不同。第一种情况中B和C都需沿第5轴与A相乘,顺序无关紧要;而后续场景中B需沿第2轴与A相乘,我更改了乘法顺序,原本未考虑性能影响,曾怀疑性能损失来自B[:,None,None,None]写法。
内容的提问来源于stack exchange,提问作者Sebastian
相关产品推荐
相关产品推荐

