Numpy数値計算コードの最適化Tips

/ Python Numpy Numerical

誰よりも速くNumpy数値計算のコードを最適化しろ！

  • Python 3.8.2
  • numpy 1.18.2
  • ipython 7.13.0 （計測用）
import numpy as np

とする。速さとコードの読みづらさはトレードオフだが仕方ない。

  • forループの除去 → Numpy関数に置き換え
  • メモリ割り当ての除去 → 明示的にバッファを用意する
  • 変数のコピーを作らない → 参照（view）を使う
  • マニュアルを読む
  • （おまけ）最新バージョン&MKLを使う

に限る。

forループを除去する

最も効果がある。ケースにもよるが100回程度のループ1個のネスト解消につき50～100倍ほど違う。ndarrayの要素に対するループ処理はnumpyで全て用意されている（といっても過言でないくらい多くのことが出来る）のでもしあれば排除する。

多くの場合にPythonの重いforを用いているのは計算手順がndarrayの計算で出来なそうだからという理由だと思うが、インデックスが添え字を回っている場合はほぼ確実にfor無しでも実現できるはず。

逐次計算を展開する

以下はいずれも平均が結果として得られる手続きになる。nで割らずにnp.sumを使えば和が得られる。

  • 0で初期化したサイズ1のバッファにforで要素を逐一加算して最後にnで割る
  • サイズnのバッファにfor以外で全要素を格納してからnp.meanを使う

後者の方がメモリは喰うがメモリが物理的に足りるなら間違いなく後者がいい。C/C++の数値計算プログラムとPython&Numpyの数値計算プログラムはforループの遅さという点で根本的に違うのでメモリ容量が最適でないとかは諦める。メモリどうこうより処理が速く返ってくる方がほとんどの場合幸せなはず。

axisキーワードを使う

np.mean()やnp.sum()がこれに当たる。多次元配列の特定の軸にそって平均や和を取る操作（行列でいう行和や列和、行平均など）はaxisキーワードを使えば実現できる。知らずにforでやっているとしたら必ず置き換えた方がいい。下記のbroadcastingのために要素数1の次元を残したい場合はkeepdims=Trueを指定するようにする。

broadcastを理解する

Numpyにはndarrayの形状を柔軟に解釈してくれるbroadcastという仕組みがあります。柔軟がゆえに人が書いたコードを読めない原因になっていたりもしますが、どっちにしろこれを理解していないと不必要にコードが長く（＝パフォーマンス低下の可能性）なったりndarrayに関するエラーメッセージの言っていることが全く分かりません。ValueError: operands could not be broadcast togetherは本当ならば柔軟に解釈したいんだけでも出来ませんという意味。

Broadcasting ― Numpy v1.17 Manual
https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html

General Broadcasting Rulesにある例を見るのが一番手っ取り早い。Broadcastingでは軸の要素数は大きい方に合わせられるので要素数と軸の数が減ることがない。

expand_dimsを使う

先述のリンク先の通りboradcastingにおいて要素数1の軸は特別な意味を持っています。ndarrayの変形は2次元であればreshapeとかで手におえますがbroadcastを前提とする多次元配列の操作だと使える時がある。特に人間は3次元までしかイメージできないのでndimを大きくしないと計算できないときなど。

A.shape : (n, N)
B.shape : (m, N)
----------------
A = np.expand_dims(A, axis=1): (n, 1, N)
B = np.expand_dims(B, axis=0): (1, m, N)
----------------
C = A*B : (n, m, N)

名前にexpandとあるだけあって列ベクトルと行ベクトルをこの順番で掛け算をしてベクトルを行列にexpandする操作に向いている。この場合は第3軸がデータ番号になっているのでそれを保持したまま比較的すっきりとこの操作を書ける。データ番号が2次元配列だった場合なんかにも動じることがない。

swapaxesとreshape

forを避けるためにはもともとこれらを使わなくて済むようにデータの入れ方を工夫するのが先と思われるが一応これらもforの除去につながる。先ほどの例の続きで今度はN個のn, m行列を別なN個のm, k行列とデータごとに行列積を取りたいとする。さらに、最終的にはこの行列を直列化したいとする（行列ごとにflattenもしくはravel）。

C.shape : (n, m, N)
D.shape : (m, k, N)
-------------------
C = np.swapaxes(C, 0, 2) : (N, m, n)
D = np.swapaxes(D, 1, 2) : (N, k, m)
-------------------
E = np.matmul(D, C)         : (N, k, n)
E = np.swapaxes(E, 1, 2) : (N, n, k)
-------------------
G = E.reshape(N, -1)     : (N, n*k)

というようになる。np.matmul()は3軸以上同士のndarrayの計算は最後2つの軸を行列と見なして行列積を計算する。np.reshape()に－1を指定すると残りの要素を全てフラットにするという意味になる。局所的にはnp.flatten()やnp.ravel()と同じだがこれらには不可能な操作である。

whereを使う

条件付き演算をndarrayの全要素に対して施す場合などはnp.where()を使ってループ無しに実現できる。要素がゼロ以下のものをゼロにするなどといった操作は以下で出来る。

X = np.random.randn(10, 10, 10)
X = np.where(X > 0, X, 0.)

使い方はExcelのIF関数に近い。こういう場合に3重ループを回してはいけない。やむを得ない場合はfor i, j, k in np.ndindex(10, 10, 10)がクールと思われる。

メモリ割り当て関連の最適化

  • マニュアルを読む
  • キャッシュを期待せずに明示的にバッファを取る
  • そもそもメモリ割り当て（コピーの作成）が発生しないような書き方

が基本。まずは使っているmethod全ての公式マニュアルを参照するところから始める。site:scipy.org指定でノイズを除去しつつGoogle検索することをオススメする。次にサイズが変化しない計算結果の格納変数や重要な途中結果を格納するような、かつ大容量の配列は半明示的に初回だけnp.zeros()などでアロケートして明示的に使いまわすようにする。クリティカルな部分では等号による代入などは避けてそもそも変数のコピーが発生しないように書く。

flattenではなくravelを使う

マニュアルで言われている通りflaten()が毎回配列のコピーを作る（Return a copy of array）のに対してravel()は必要な時にしかコピーを作らない（A copy is made only if needed）。

In [0]: x = np.random.randn(100, 100)

In [1]: %timeit x.flatten()
3.47 µs ± 100 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [2]: %timeit x.ravel()
243 ns ± 7.82 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

flatten()を使わなければならない状況はあまり思いつかない。ravel()を使うべき。

引数にoutがあれば使う

デフォルト引数としてout=Noneをもっているものがある。

In [0]: A, x, y = np.random.randn(100, 100), np.random.randn(100, 1), np.zeros((100, 1))

In [1]: %timeit y[:,:] = np.dot(A,x)
3.75 µs ± 22.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [2]: %timeit np.dot(A,x,out=y)
2.88 µs ± 51.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

見ての通りこれをつかうことそのものによる効果は限定的。重要なのは毎回結果のコピー先がアロケートされないように注意すること。上記2例は2つともそれを満たしているので両方速いほう。y = np.dot(A,x)はコピーの可能性のある書き方なのと前のyのガーベジが溜っていく書き方なので避ける。

ゼロ埋めのためにzerosを使わない

np.zerosは見た目上配列がアロケートされる可能性があるのでもし既にバッファがあるならndarray.fill(0)を使う。

In [0]: x = np.zeros(100000)

In [1]: %timeit x = np.zeros(100000)
25.4 µs ± 324 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [2]: %timeit x.fill(0)
20.5 µs ± 366 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

こちらも効果は限定的だが簡単でかつ着実な効果がある。np.ones(n)*10のような使い方で全要素を10で埋めるときも既にバッファがあるなら確実にndarray.fill(10)のほうがよい（このケースの方が効果は大きそうである）。

おまけ：外部的な要因

以下は気分的なもの。

  • PythonとNumpyは最新バージョンに保っておく
  • Intel Math Kernel Library（mkl）でビルドされたものを使う
In [0]: np.show_config()
blas_mkl_info:
    libraries = ['mkl_lapack95_lp64', 'mkl_blas95_lp64', 'mkl_rt']
    library_dirs = ** 略 **
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ** 略 **
blas_opt_info:
    libraries = ['mkl_lapack95_lp64', 'mkl_blas95_lp64', 'mkl_rt']
    library_dirs = ** 略 **
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ** 略 **
lapack_mkl_info:
    libraries = ['mkl_lapack95_lp64', 'mkl_blas95_lp64', 'mkl_rt']
    library_dirs = ** 略 **
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ** 略 **
lapack_opt_info:
    libraries = ['mkl_lapack95_lp64', 'mkl_blas95_lp64', 'mkl_rt']
    library_dirs = ** 略 **
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ** 略 **

みたいにnp.show_config()でmklの文字が散見されればMKLによってビルドされている（はず）。

リンク集

以下から何か役立ちそうなものが見つかったらメモしていく。

Routines ― Numpy v1.17
https://docs.scipy.org/doc/numpy/reference/routines.html