You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

如何避免轴交换或虚拟维度提升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

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.06.23 23:57:18