为什么用 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 语法更清晰（也更快）——从视觉上看，它作为一个维度添加操作脱颖而出。

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

Why isn't broadcasting with numpy faster than a nested loop

python

arrays

numpy

nested-loops

array-broadcasting

时间