为什么用 numpy 广播的速度不比嵌套循环快

Question

我的代码中有一个计算被执行了数千次，我想看看我是否可以让它更快，因为它目前使用两个嵌套循环。我认为如果我使用广播，我可以让它快几倍。

我在下面展示了两个选项，谢天谢地，它们给出了相同的结果。

import numpy as np

n = 1000
x = np.random.random([n, 3])
y = np.random.random([n, 3])
func_weight = np.random.random(n)


result = np.zeros([n, 9])
result_2 = np.zeros([n, 9])

# existing
for a in range(3):
    for b in range(3):
        result[:, 3*a + b] = x[:, a] * y[:, b] * func_weight

# broadcasting - assumed this would be faster
for a in range(3):
    result_2[:, 3*a:3*(a+1)] = np.expand_dims(x[:, a], axis=-1) * y * np.expand_dims(func_weight, axis=-1)

时间

n=100
nested loops: 24.7 µs ± 362 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
broadcasting: 70.3 µs ± 1.22 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

n=1000
nested loops: 50.5 µs ± 913 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
broadcasting: 148 µs ± 372 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

n=10000
nested loops: 327 µs ± 7.99 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
broadcasting: 864 µs ± 5.57 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

在我的测试中，广播总是比较慢，所以我对正在发生的事情有点困惑。我猜是因为我必须使用expand_dims 来使第二个解决方案中的形状对齐，这就是对性能的重大影响。那是对的吗？随着数组大小的增加，嵌套循环的性能并没有太大变化，总是快 3 倍左右。

有没有我没有考虑过的更优化的第三种解决方案？

Answer 1

In [126]: %%timeit
     ...: result = np.zeros([n,9])
     ...: for a in range(3):
     ...:     for b in range(3):
     ...:         result[:, 3*a + b] = x[:, a] * y[:, b] * func_weight
141 µs ± 255 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [128]: %%timeit
     ...: result_2 = np.zeros([n,9])
     ...: for a in range(3):
     ...:    result_2[:, 3*a:3*(a+1)] = np.expand_dims(x[:, a], axis=-1) * y * n
     ...: p.expand_dims(func_weight, axis=-1)
202 µs ± 10.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

一个完整的广播版本：

In [130]: %%timeit
     ...: result_3 = (x[:,:,None]*y[:,None,:]*func_weight[:,None,None]).reshape(
     ...: n,9) 
88.8 µs ± 73.1 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

将expand_dims 替换为np.newaxis/None 扩展：

In [131]: %%timeit
     ...: result_2 = np.zeros([n,9])
     ...: for a in range(3):
     ...:    result_2[:, 3*a:3*(a+1)] = x[:, a,None] * y * func_weight[:,None]
132 µs ± 315 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

所以是的，expand_dims 有点慢，我认为是因为它试图成为通用的。还有一层额外的函数调用。

expand_dims 只是a.reshape(shape)，但是将您的轴参数转换为shape 元组需要一些时间。作为一个有经验的用户，我发现None 语法更清晰（也更快）——在视觉上它作为一个维度添加动作脱颖而出。