聚合 numpy 函数

Question

我有一个 numpy 操作，我需要密集调用，我需要优化：

 np.sum(a**2, axis=1)**.5   # where a is a 2 dimensional ndarray

这个操作由三个函数组成，需要迭代 'a' 3 次。将所有操作聚合在一个函数下并沿轴 1 仅应用该函数一次会更有效。不幸的是，numpy 的 apply_along_axis 函数不是一个选项，因为性能大约差 1000 倍。

有没有办法聚合几个 numpy 操作，所以它只需要在数组上循环一次？

Answer 1

使用浮点数组时，可以使用np.einsum -

np.sqrt(np.einsum('ij,ij->i',a,a))

运行时测试-

In [34]: a = np.random.rand(1000,1000)

In [35]: np.allclose(np.sum(a**2, axis=1)**.5,np.sqrt(np.einsum('ij,ij->i',a,a)))
Out[35]: True

In [36]: %timeit np.sum(a**2, axis=1)**.5
100 loops, best of 3: 7.57 ms per loop

In [37]: %timeit np.sqrt(np.einsum('ij,ij->i',a,a))
1000 loops, best of 3: 1.52 ms per loop

Answer 2

看看numexpr，它可以让你比纯numpy更快地计算数值表达式：

In [19]: a = np.arange(1e6).reshape(1000,1000)

In [20]: import numexpr as ne

In [21]: %timeit np.sum(a**2,axis=1)**0.5
100 loops, best of 3: 6.08 ms per loop

In [22]: %timeit ne.evaluate("sum(a**2,axis=1)")**0.5
100 loops, best of 3: 4.27 ms per loop

**0.5 不是表达式的一部分，因为sum 是归约运算，需要在表达式中最后计算。您还可以对sqrt/**0.5 运行另一个评估。