类似代码之间numpy的巨大速度差异

Question

为什么下面的L2范数计算会有这么大的速度差异：

a = np.arange(1200.0).reshape((-1,3))

%timeit [np.sqrt((a*a).sum(axis=1))]
100000 loops, best of 3: 12 µs per loop

%timeit [np.sqrt(np.dot(x,x)) for x in a]
1000 loops, best of 3: 814 µs per loop

%timeit [np.linalg.norm(x) for x in a]
100 loops, best of 3: 2 ms per loop

据我所知，所有三个都产生相同的结果。

这里是 numpy.linalg.norm 函数的源代码：

x = asarray(x)

# Check the default case first and handle it immediately.
if ord is None and axis is None:
    x = x.ravel(order='K')
    if isComplexType(x.dtype.type):
        sqnorm = dot(x.real, x.real) + dot(x.imag, x.imag)
    else:
        sqnorm = dot(x, x)
    return sqrt(sqnorm)

编辑：有人建议可以并行化一个版本，但我检查了一下，事实并非如此。所有三个版本都消耗 12.5% 的 CPU（这通常是我的 4 个物理/8 个虚拟核 Xeon CPU 上的 Python 代码的情况。

Answer 1

np.dot 通常会调用 BLAS 库函数 - 因此它的速度将取决于您的 numpy 版本链接到的 BLAS 库。一般来说，我希望它具有更大的恒定开销，但随着数组大小的增加，可以更好地扩展。但是，您从列表解析中调用它（实际上是一个普通的 Python for 循环）这一事实可能会抵消使用 BLAS 的任何性能优势。

如果您摆脱列表理解并使用 axis= kwarg，np.linalg.norm 与您的第一个示例相当，但 np.einsum 比两者都快得多：

In [1]: %timeit np.sqrt((a*a).sum(axis=1))
The slowest run took 10.12 times longer than the fastest. This could mean that an intermediate result is being cached 
100000 loops, best of 3: 11.1 µs per loop

In [2]: %timeit np.linalg.norm(a, axis=1)
The slowest run took 14.63 times longer than the fastest. This could mean that an intermediate result is being cached 
100000 loops, best of 3: 13.5 µs per loop

# this is what np.linalg.norm does internally
In [3]: %timeit np.sqrt(np.add.reduce(a * a, axis=1))
The slowest run took 34.05 times longer than the fastest. This could mean that an intermediate result is being cached 
100000 loops, best of 3: 10.7 µs per loop

In [4]: %timeit np.sqrt(np.einsum('ij,ij->i',a,a))
The slowest run took 5.55 times longer than the fastest. This could mean that an intermediate result is being cached 
100000 loops, best of 3: 5.42 µs per loop