为什么用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)，但是将axis参数转换为shape元组需要一些时间。作为一个有经验的用户，我发现None语法更清晰（更快）-从视觉上看，它是一个维度添加操作