python - numpy np.apply_along_axis 函数加速？

Question

np.apply_along_axis() 函数似乎非常慢(15 分钟后没有输出)。有没有一种快速的方法可以在长数组上执行此功能而无需并行化操作？我专门讨论的是具有数百万个元素的数组。

这是我正在尝试做的一个例子。请忽略 my_func 的简单定义，目标不是将数组乘以 55(当然无论如何都可以就地完成)，而是一个说明。在实践中，my_func 稍微复杂一些，需要额外的参数，因此 a 的每个元素都被不同地修改，即不仅仅是乘以 55。

>>> def my_func(a):
...     return a[0]*55
>>> a = np.ones((200000000,1))
>>> np.apply_along_axis(my_func, 1, a)

编辑:

a = np.ones((20,1))

def my_func(a, i,j):
...     b = np.zeros((2,2))
...     b[0,0] = a[i]
...     b[1,0] = a[i]
...     b[0,1] = a[i]
...     b[1,1] = a[j]
...     return  linalg.eigh(b)


>>> my_func(a,1,1)
(array([ 0.,  2.]), array([[-0.70710678,  0.70710678],
   [ 0.70710678,  0.70710678]]))

Answer 1

np.apply_along_axis 不是为了速度。

没有办法将纯 Python 函数应用到 Numpy 数组的每个元素而不调用它多次，除非 AST 重写...

还好有解决办法:

• 矢量化

  虽然这通常很难，但通常是简单的解决方案。找到某种方法以概括元素的方式来表达您的计算，这样您就可以一次处理整个矩阵。这将导致循环被提升出 Python 并进入经过高度优化的 C 和 Fortran 例程。

• JITing:Numba 和 Parakeet，以及较小程度的 PyPy 和 NumPy

  Numba 和 Parakeet 都处理 Numpy 数据结构上的 JITing 循环，因此如果您将循环 内联到 一个函数(这可以是一个包装函数)，您几乎可以获得巨大的速度提升 -自由的。不过，这取决于所使用的数据结构。

• 符号评估器，例如 Theano 和 numexpr

  这些允许您使用嵌入式语言来表达计算，这甚至可以比矢量化版本快得多。

• Cython 和 C 扩展

  如果所有其他内容都丢失了，您始终可以手动挖掘到 C。Cython 隐藏了很多复杂性，并且也有很多可爱的魔法，所以它并不总是那么糟糕(尽管它有助于了解您是什么做)。

---

给你。

这是我的测试“环境”(你真的应该提供这个:P):

import itertools
import numpy

a = numpy.arange(200).reshape((200,1)) ** 2

def my_func(a, i,j):
    b = numpy.zeros((2,2))
    b[0,0] = a[i]
    b[1,0] = a[i]
    b[0,1] = a[i]
    b[1,1] = a[j]
    return  numpy.linalg.eigh(b)

eigvals = {}
eigvecs = {}

for i, j in itertools.combinations(range(a.size), 2):
    eigvals[i, j], eigvecs[i, j] = my_func(a,i,j)

现在，获得所有排列而不是组合要容易得多，因为您可以这样做:

# All *permutations*, not combinations
indexes = numpy.mgrid[:a.size, :a.size]

这可能看起来很浪费，但排列只有两倍，所以没什么大不了的。

所以我们要使用这些索引来获取相关元素:

# Remove the extra dimension; it's not wanted here!
subs = a[:,0][indexes]

然后我们可以制作我们的矩阵:

target = numpy.array([
    [subs[0], subs[0]],
    [subs[0], subs[1]]
])

我们需要矩阵在最后两个维度:

target.shape
#>>> (2, 2, 200, 200)

target = numpy.swapaxes(target, 0, 2)
target = numpy.swapaxes(target, 1, 3)

target.shape
#>>> (200, 200, 2, 2)

我们可以检查它是否有效:

target[10, 20]
#>>> array([[100, 100],
#>>>        [100, 400]])

耶!

那么我们就运行 numpy.linalg.eigh:

values, vectors = numpy.linalg.eigh(target)

看，它有效!

values[10, 20]
#>>> array([  69.72243623,  430.27756377])

eigvals[10, 20]
#>>> array([  69.72243623,  430.27756377])

那么我想你可能想要连接这些:

numpy.concatenate([values[row, row+1:] for row in range(len(values))])
#>>> array([[  0.00000000e+00,   1.00000000e+00],
#>>>        [  0.00000000e+00,   4.00000000e+00],
#>>>        [  0.00000000e+00,   9.00000000e+00],
#>>>        ..., 
#>>>        [  1.96997462e+02,   7.78160025e+04],
#>>>        [  3.93979696e+02,   7.80160203e+04],
#>>>        [  1.97997475e+02,   7.86070025e+04]])

numpy.concatenate([vectors[row, row+1:] for row in range(len(vectors))])
#>>> array([[[ 1.        ,  0.        ],
#>>>         [ 0.        ,  1.        ]],
#>>> 
#>>>        [[ 1.        ,  0.        ],
#>>>         [ 0.        ,  1.        ]],
#>>> 
#>>>        [[ 1.        ,  0.        ],
#>>>         [ 0.        ,  1.        ]],
#>>> 
#>>>        ..., 
#>>>        [[-0.70890372,  0.70530527],
#>>>         [ 0.70530527,  0.70890372]],
#>>> 
#>>>        [[-0.71070503,  0.70349013],
#>>>         [ 0.70349013,  0.71070503]],
#>>> 
#>>>        [[-0.70889463,  0.7053144 ],
#>>>         [ 0.7053144 ,  0.70889463]]])

也可以在 numpy.mgrid 之后执行此连接循环以将工作量减半:

# All *permutations*, not combinations
indexes = numpy.mgrid[:a.size, :a.size]

# Convert to all *combinations* and reduce the dimensionality
indexes = numpy.concatenate([indexes[:, row, row+1:] for row in range(indexes.shape[1])], axis=1)

# Remove the extra dimension; it's not wanted here!
subs = a[:,0][indexes]

target = numpy.array([
    [subs[0], subs[0]],
    [subs[0], subs[1]]
])

target = numpy.rollaxis(target, 2)

values, vectors = numpy.linalg.eigh(target)

是的，最后一个样本就是你所需要的。