Pythonと関数単位での高速化(cython, numba, swig)

June 25, 2020 by Chachay
,

Python、計算が遅いです。javascriptよりも遅いです。

しかし、Python as グルー言語として「データの読み込み」や「計算処理」、「可視化」を橋渡しする機能に注目すると本当に便利なプログラミング言語だと思います。 橋渡し先の「計算処理」の定番としてはnumpyがあったり、scipyがあったりしますが、 ここをオリジナルの計算処理に置き換える手段として、 pybind11, swig, numbaやcythonがあると思っています。 使い勝手や速度の面で気になっていたので比較しました。

環境

  • Windows 10
  • Visual Studio 2019
  • miniconda - python 3.7
    • swig
    • numba
    • cython

やったこと

LU分解がベンチマークでは良く使われるそうですが、実装を簡単にするため、コレスキー分解を選びました。 ベンチマークではランダムに生成した実対称正方行列(の正定行列)を使っています。

scipyの例をあげると

# 128 x 128の行列を生成(numpy)
A128 =  gen_sym_matrix(128)

from scipy.linalg import cholesky
# コレスキー分解の時間を計測
%timeit l = cholesky(A128, lower=True)

結果

リポジトリ

python_bench/cholesky at master · Chachay/python_bench

実行時間

行列のサイズで2水準計測。

python_bench/main.ipynb at master · Chachay/python_bench

  方法 128x128
[msec]		 256x256
[msec]		 コメント
1 scipy 0.09 0.37 内部でLAPACKを呼び出し. 参考値.
2 Swig - CPP - Eigen 0.17 0.91 ホットスポットをCPPで書きたい場合、バランスが良い
インターフェース・ファイル(.i)さえ用意すれば、cppっぽく書ける。 コンパイラの差し替え(Clangとか)で更に性能が伸ばせるかも。
3 Numba - numpy 0.35 1.79 ほとんど、そのままpythonなのに魔術的な力を感じる。Pythonでプロトタイプした直後に処理速度が気になったときの第一候補.
4 Swig - CPP 0.25 2.27 スタンダードC++のみ。もっと速くできるかも。
5 Cython - Numpy 11.6 45.8 お前にはいつもがっかりしている. 俺のことが嫌いなのか。
6 Python - Numpy 21.8 87 いわゆる標準的なアプローチ
7 javascript - 295 飛び入り参加。Pythonとは関係ない。
8 Pure Python 44.4 327 List in Listで行列を表現し、組み込み関数のみで作成。
List in List形式はnumbaやcythonで上手に取り扱えないので、高速化はなし。

実行メモ

Numba

Pythonのnumpy版関数にjitデコレータをつけただけ。numpyの定義でdtypeなどで型を指定してやる必要がある。引数やローカル変数は型推論でうまく処理したのかな。

from numba import jit
@jit(nopython=True)
def cholesky_numpy_numba(A):
    M = A.shape[0]
    L = np.zeros_like(A, dtype=np.float64)

    for i in range(M):
        for j in range(0, i+1):
            s = np.dot(L[i,:j],L[j,:j])
            L[i, j] = np.sqrt(A[i,i]-s) if(i==j) else ((A[i,j]-s)/L[j,j])
    return L

Swig - Numpy

公式資料: numpy.i: a SWIG Interface File for NumPy — NumPy v1.18 Manualにいろいろ詳しい。闇についても若干言及がある。

ファイル構成

swig
├── numpy.i
├── setup.py
├── swig_mod.cpp
├── swig_mod.h
└── swig_mod.i

setup.py

ビルドの条件を書きます。モジュール名の頭にアンダースコア(_)をつけないとリンカが悲鳴をあげます。これは超重要です。

コンパイラの最適化オプションとかコンパイラの差し替えとかもここで定義できた気がします。思い出したら追記します。

SwigMod = Extension("_swig_eigen_mod", # ここのモジュール名にアンダースコア
                  ["swig_mod.i","swig_mod.cpp"],
                  swig_opts=['-c++', '-py3'],
                  include_dirs = [numpy_include]
                  )

swig_mod.cpp & swig_mod.h

本体。Swigのマクロで2次元配列のnumpyの配列を(行サイズ、列サイズ、ポインタ)に置き換えるので、

import numpy as np
from SwigMod import cholesky_swig
A = np.zeros((3,3))
L = cholesky_swig(A)

のような使い方を予定しているモジュールを作るときには、引数をint mm1, int mn1, double* mat1のように定義しておき、 戻り値も引数としてdouble** outmat, int* mm, int* mnのように定義しておく。 

// 関数の定義としては戻り値はvoid
void cholesky_swig(
          // 引数Aの成れの果て
          int mm1, int mn1, double* mat1,
          // 返り値Lの素
          double** outmat, int* mm, int* mn){
  // メモリの確保はC++で行い、そのまま管理をpythonに引き継ぐ
  double* arr = (double*)calloc(mn1 * mm1, sizeof(double));

  // ---- 中略 ------

  // 返り値はポインタで渡す。returnしない。
  *mm = mm1;
  *mn = mn1;
  *outmat = arr;
}

戻り値についてはSwigがラッパーを作るときに、引数を戻り値に置き換えるらしい。黒魔術なので深入りしてない。公式資料に詳しい。numpyの1次元配列、多次元配列についても同様。

後述のSwig - Eigenのようにnumpyをtypedefして、きちんとインターフェースを書くとこんな気持ち悪いことにはならないのかもしれない。

callocで確保したヒープが最後どのようになるのか杳として知れず。

swig_mod.i

ヘッダーファイルとpythonをつなぐラッパーを生成するためのスクリプト。numpyを開くところについては下記でマクロ定義。変数名はヘッダーで使ったものと一致するように。

%apply (double** ARGOUTVIEWM_ARRAY2, int* DIM1, int* DIM2){(double** outmat, int* mm, int* mn)}
%apply (int DIM1, int DIM2, double* IN_ARRAY2){(int mm1, int mn1, double* mat1)}
// 複数定義する場合
// %apply (int DIM1, int DIM2, double* IN_ARRAY2){(int mm1, int mn1, double* mat1), (int mm2, int mn2, double* mat2)}

モジュールのビルド

python setup.py build_ext --inplace

するとsetup.pyと同じ階層にモジュールが作られるので、jupyterから呼び出しやすいところへコピーして使う。

Swig - Eigen

ありがたいベストプラクティスに出会えたので追加。

eigen.iが秘伝のタレです。numpyとeigenを橋渡ししてくれます。46行目の#include "Eigen/Core"だけ環境によっては書き換え必要。

インターフェースファイルの書き方も詳しいので熟読するべき

swig-eigen-numpy/inverter_wrapper.i

eigen.iがしっかりしているので、hppもこれがpythonモジュールの定義ファイルなのか疑うレベルできれいなC++になります。

#include <exception>
#include <vector>
#include "eigen/Eigen/Core"

Eigen::MatrixXd cholesky_swig_eigen(const Eigen::MatrixXd &M);

Cython

お前にはがっかりしている。どうしたら良いか私も知りたい。

%%cython
import numpy as np
cimport numpy as np
cimport cython

ctypedef double DTYPE_t

@cython.boundscheck(False)
@cython.wraparound(False)
def cholesky_numpy_cython(np.ndarray[DTYPE_t, ndim=2]  A):
    cdef int i
    cdef int j
    cdef double s
    cdef int M = A.shape[0]
    cdef np.ndarray[DTYPE_t, ndim=2] L = np.zeros_like(A, dtype=np.float64)

    for i in range(M):
        for j in range(0, i+1):
            s = np.dot(L[i,:j],L[j,:j])
            L[i, j] = (A[i,i]-s)**0.5 if(i==j) else ((A[i,j]-s)/L[j,j])
    return L

参考文献

WindowsにSWIG導入 - Volo di notte
いまやconda install swigで楽にインストールできるのですが、anacondaを使いたくない場合などに参照ください
Write C++ extensions for Python - Visual Studio | Microsoft Docs
今後よみたい
SWIG and Python
ちょっと古い話もある
at

No comments:

Post a Comment