numpyを使って書かれたpythonプログラムをコンパイルしてから実行することで高速化するモジュール.一般的にpythonのfor文なんかを使って書かれているコードはpythonオブジェクトを使って走るのですこぶる遅いが,numbaのJIT(Just-In-Time)コンパイル機能を使ってコンパイルして走らせることで100倍とか速くすることが可能.pythonプログラムはC++とかで書かれたコードに比べて1000倍とか遅いことが多いが,JIT化することでC++に比べて2~5倍程度遅いところまで肉薄できるようだ.しかも,少し文言を追加するだけで,コードの改変量はさほど多くない. 科学技術計算や機械学習では重たい計算もあるので,numbaの存在を覚えておくと良い.

インストール

$ pip install -U numba

簡単な例

ある関数 func を高速化させたい場合,次のように書く.

from numba import njit, prange
 
@njit("i4(i4,i4)",parallel=True)
def func(istart,iend):
    sum = 0
    for i in prange(istart,iend):
        sum += i
    return sum
  • 関数前のデコレータで関数funcをJIT化するという宣言
  • njitjit(nopython=True) の略で,pythonオブジェクトコードを使う場合にエラーで止まることを意味している.pythonオブジェクトを使うとコンパイルしても速くならないので意味ないからエラーで止まってくれた方が良い.
  • 引数として,"i4(i4,i4)" のように関数の返り値や引数の型を明示的に指定(しなくても勝手に判断してくれると思われるが.)
  • 変数の型は実際には int32 だとか float64 だとか長い名前があるようだが,短縮形も定義されているからnumbaのウェブサイト(例えばココ)で確認すると良い.
  • parallel=True とするとこでスレッド並列化までしてくれる.prangeという関数はその並列化をどのループで行うか指定している.

上のようにJIT化された関数は最初に呼び出したとき(?)にコンパイルされて,コンパイル済みコードを実行することになるからだいぶ高速になる.

他にも,下記のようにしてJIT化された関数を定義することもできる.

func_jit = njit(func)

注意事項

  • 関数によってはJIT化できないものもある.JIT化したい関数の中に独自の関数なんかがあってもJIT化できないと思われる.その場合は,その関数もJIT化すればよいのかな?
  • たとえば,PyTorchの関数なんかがあるとJIT化できない.(PyTorchは独自のJIT機構を持っているようだが)