この記事はMicroAd (マイクロアド) Advent Calendar 2020の18日目の記事です。

はじめに

numbaとはpythonコードを高速化するためのライブラリです。

高速化したい関数に以下のようにアノテーションをつけるだけでネイティブのマシンコードのような速度で実行できるようになります。

import numba

@numba.jit
def f(x, y):
    return x + y

numbaを使うとお手軽にコードを高速化できる反面、通常のpythonコードとは違いコンパイルが発生するので
仕組みを理解せずに使っていると返って処理全体の実行速度が遅くなってしまうこともあります。

この記事では意外なnumbaの落とし穴とそれを回避するための方法をご紹介しようと思います。

no pythonモードを利用する

numbaには最適化が行われるno pythonと通常と同じコンパイル方式のobjectの２種類のモードが存在します。

この二つのモードはデフォルト設定では自動で切り替わるようになっており、
no pythonモードが優先されますが、これが失敗した場合は自動でobjectモードに切り替わります。

オブジェクトモードでは最適化が行われない上、
コンパイルのオーバーヘッドでかえって実行時間がかかる場合があります。

なのでnumbaを活用するには最適化に対応していない処理を把握しておき、
そういった処理を含まないよう関数を切り出すなどの工夫が重要です。

no pythonモード非対応の処理
・dict
・set
・pandas
・yield from
・set・dict・ジェネレーターの内包表記

また、numbaではno pythonモードでのコンパイルができない場合にエラーをですように設定できるので、
基本的には以下の書き方で利用するのをお勧めします。

@jit(nopython=True)
def f(x, y):

@njit
def f(x, y):

コンパイルのオーバーヘッドを回避する

numbaでは実装に応じて関数をコンパイルするタイミングが自動で切り替わるようになっています。

それを把握せずに利用すると処理速度が求められる場面でコンパイルが走ってしまい、
高速化のために利用したnumbaが返ってボトルネックになってしまうかもしれません。

意識的にコンパイルタイミングを操作するため、そのバリエーションと書き方を把握しておきましょう。

関数が初めて呼ばれたタイミングでコンパイル

こちらがデフォルトの動作です。

@njit
def f(x, y):
  return x + y

関数が読み込まれたタイミングでコンパイル

型指定をすることで、関数が定義されたタイミングでコンパイルされるようになります。
型指定には高速化の効果もあるので、この書き方が一番のおすすめです。

@njit(int32(int32, int32))
def f(x, y):
    return x + y

cacheを利用する

またcacheを利用するよう設定することで、初回起動以外はコンパイルを必要なくできます。
この方法は他の型指定とも併用できるので、とりあえずでオプションをつけておいても損はないです。

@njit(cache=True)
def f(x, y):
  return x + y

事前コンパイル

コンパイル済みモジュールの自動作成することでコンパイルを実行時に起きなくすることもできます。
正直この方法はちょっと面倒くさいですしコードも冗長になります。
どんなタイミングでもコンパイルによるオーバーヘッドを起こしたくない時の最終手段ですね。

from numba.pycc import CC

cc = CC('my_module')

@cc.export('multi', 'i4(i4, i4)')
def f(a, b):
    return a + b

if __name__ == "__main__":
    cc.compile()
>>> from my_module import f
>>> f(3, 4)

まとめ

いかがだったでしょうか。
今回は簡単に高速化を実現できますが、ちょっと動作に癖があるnumbaの注意点をまとめてみました。
この記事が読んでいただいた皆さんの参考になれば幸いです。

numbaでpythonコードを正しく高速化するためのTips