快速（呃）numpy花式索引和减少？

Posted 2023-03-05

【中文标题】快速（呃）numpy花式索引和减少？

【英文标题】：Fast(er) numpy fancy indexing and reduction?

【发布时间】：2012-08-01 18:13:10

【问题描述】：

我正在尝试使用和加速花式索引来“加入”两个数组并在结果轴之一上求和。

类似这样的：

$ ipython
In [1]: import numpy as np
In [2]: ne, ds = 12, 6
In [3]: i = np.random.randn(ne, ds).astype('float32')
In [4]: t = np.random.randint(0, ds, size=(1e5, ne)).astype('uint8')

In [5]: %timeit i[np.arange(ne), t].sum(-1)
10 loops, best of 3: 44 ms per loop

有没有一种简单的方法可以加速In [5] 中的语句？我应该使用 OpenMP 和 scipy.weave 或 Cython 的 prange 之类的吗？

【问题讨论】：

另一个相关问题是我如何使用pandas 做同样的事情？

Numpy 以 C 速度执行此操作，因此您可能无法通过 weave 加快速度。

【参考方案1】：

numpy.take 由于某种原因比花哨的索引要快得多。唯一的技巧是它将数组视为平面。

In [1]: a = np.random.randn(12,6).astype(np.float32)

In [2]: c = np.random.randint(0,6,size=(1e5,12)).astype(np.uint8)

In [3]: r = np.arange(12)

In [4]: %timeit a[r,c].sum(-1)
10 loops, best of 3: 46.7 ms per loop

In [5]: rr, cc = np.broadcast_arrays(r,c)

In [6]: flat_index = rr*a.shape[1] + cc

In [7]: %timeit a.take(flat_index).sum(-1)
100 loops, best of 3: 5.5 ms per loop

In [8]: (a.take(flat_index).sum(-1) == a[r,c].sum(-1)).all()
Out[8]: True

我认为，除此之外，您会看到速度大幅提升的唯一其他方法是使用 PyCUDA 之类的东西为 GPU 编写自定义内核。

【讨论】：

它默认只将数组视为平面，您仍然可以使用axis关键字。即np.take(np.arange(10).reshape((-1,2)), [0], axis=0) 将选择第一行。

@jorgeca：是的，但我认为您不能像使用花哨的索引那样通过指定行和列来提取单个元素，除非您索引平面数组。