Question

想象一个形状为N，C，K，H，W的Numpy np.array
和 meanCoeffs = np.mean(x, axis=(3,4))的形状为N，C，K

import time
import numpy as np

def func(j1,t1,L):
    return t1+j1*L


N,C,J,L,H,W = 128,5,3,8,64,64
x = np.random.rand(N,C,J*L,H,W)
meanCoeffs = np.mean(x, axis=(3,4))
xnew=np.zeros_like(x)

start = time.time()
for ib in range(x.shape[0]):
    for filt in range(x.shape[1]):
        for j1 in range(J):
            for t1 in range(L):
                i=func(j1,t1,L)
                xnew[ib,filt,i,:,:] = x[ib,filt,i,:,:]-meanCoeffs[ib,filt,i]
end = time.time()
print ("Time elapsed:", end - start)

经过的时间：0.18448996543884277

请注意，func函数可能很复杂。

有没有办法加快这种原始代码的速度？

感谢您的帮助。

Answer 1

索引中

for j1 in range(J):
    for t1 in range(L):
       i=t1+j1*L

相当于简单地说

for i in range(J*L):

也就是说，三个循环归结为所有三个索引的循环，减法是一个简单的广播：

xnew2 = x - meanCoeffs[:, :, :, None, None]

如果您愿意，也可以使用np.expand_dims(meanCoeffs, (3, 4))。

您可以验证这与您在具体示例中所做的相同：

In [14]: np.all(xnew == xnew2)
Out[14]: True

python / numpy：如何tweek嵌套循环以加快代码速度

1 个答案: