11

私は、2つの配列を「結合」し、結果の軸の1つを合計するために、ファンシーインデックスを使用して加速しようとしています。

このようなもの：

$ 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.weaveorCythonのようなものを使用する必要がありprangeますか?

4

1 に答える 1

8

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用のカスタムカーネルを作成することだと思います。

于 2012-08-05T00:27:59.793 に答える