Posted inプログラミング

Numba JITでPythonを高速化：重いデータ処理を100倍速くする方法

Programming tutorial - IT technology blog
Programming tutorial - IT technology blog

大量データでPythonスクリプトが遅くなる理由

純粋なPythonでシミュレーションを実行したり、数百万件のデータを処理しようとすると、本当に辛い。45分かかるスクリプトをじっと眺めながら、コンパイル言語なら30秒で処理できるのにと思ったことが何度もある。腹立たしいのは、ロジック自体は正しいという点だ。Pythonが遅いだけなのだ。

根本的な原因は意外なことではない。Pythonはインタープリタ言語だ。すべての行が実行時にインタープリタを通り、変数の型チェック、メソッド名の解決、メモリの動的管理が行われる。Webアプリやちょっとしたスクリプトではこのオーバーヘッドは見えない。しかし数値計算——数百万要素のループ、反復シミュレーション、条件分岐を持つ金融モデル——ではオーバーヘッドが急速に積み重なる。

NumPyが一般的な解決策として使われており、大幅な改善をもたらす。内部でコンパイル済みのCコードを呼び出すため、配列全体のベクトル化演算は高速に動作する。しかし、カスタムループ——反復アルゴリズム、物理シミュレーション、信号処理ルーティン——が必要になった途端、低速なPythonに戻ってしまう。NumPyではそこは助けられない。

そこでNumbaの出番だ。PythonのJIT（Just-In-Time）コンパイラである。関数を毎回インタープリタで実行する代わりに、Numbaは初回の呼び出し時にネイティブマシンコードにコンパイルする。それ以降の呼び出しはコンパイル済みバイナリに直接アクセスする——これが100倍の高速化をもたらす理由だ。Pythonで本格的な数値計算をしているなら、これはあなたのスタックで最も費用対効果の高いツールの一つだろう。

インストール

NumbaのハードデペンデンシーはNumPyのみで、ほとんどのPythonデータ環境にはすでに入っている。以下でインストールできる：

pip install numba

NVIDIAハードウェアでGPUアクセラレーションを使いたい場合は、先にNVIDIAのサイトからCUDA Toolkitを別途インストールし、次を追加する：

pip install numba cuda-python

このガイドのすべての内容は、基本的なpipインストールで十分だ。動作確認：

import numba
print(numba.__version__)

NumbaはPython 3.9+とNumPy 1.20+で最もよく動作する。古い環境の場合は、コードとは無関係なコンパイルエラーに悩まされる前にアップグレードしよう。

設定：Numbaを使いこなす

@jitデコレータ——最初の入口

@jitデコレータから始めよう——最もシンプルな入口だ。任意の関数に追加すると、Numbaは初回の呼び出し時にコンパイルする：

from numba import jit
import numpy as np

@jit
def sum_squares(arr):
    total = 0.0
    for i in range(len(arr)):
        total += arr[i] ** 2
    return total

data = np.random.rand(10_000_000)
result = sum_squares(data)  # 初回呼び出し：コンパイル＋実行

初回の呼び出しでコンパイルが始まる——少し間が空くことを予期しておこう。その後の呼び出しはすべてコンパイル済みマシンコードを直接実行する。

実際のパフォーマンス保証には@njitを使う

@jitには無音のフォールバックがある：Numbaが関数をコンパイルできない場合、エラーを発生させずに通常のPythonとして静かに実行される。手早い実験には便利だが、パフォーマンスの問題を隠してしまう。代わりに@njit（nopythonモード）を使おう——コンパイルできないものがあれば即座にエラーを発生させるので、何を修正すべきかが正確にわかる：

from numba import njit

@njit
def compute_distance(x1, y1, x2, y2):
    return ((x2 - x1)**2 + (y2 - y1)**2) ** 0.5

# Pythonにフォールバックした場合、Numbaはここでエラーを発生させる
print(compute_distance(0.0, 0.0, 3.0, 4.0))  # 5.0

@njitがエラーを出す場合、エラーメッセージはサポートされていない部分を指摘する——通常は文字列、通常のPythonの辞書、カスタムクラスなどの非数値型だ。

parallel=Trueを使った並列ループ

互いに依存しないイテレーションには、Numbaがprange（並列レンジ）を使って自動的にCPUコアにワークを分散できる：

from numba import njit, prange
import numpy as np

@njit(parallel=True)
def parallel_sum_squares(arr):
    total = 0.0
    for i in prange(len(arr)):  # rangeではなくprange
        total += arr[i] ** 2
    return total

data = np.random.rand(10_000_000)
result = parallel_sum_squares(data)

Numbaは利用可能なすべてのCPUコアにループを自動的に分割する。8コアマシンなら、JITの高速化に加えてさらに4〜6倍の速度向上が得られる。

コンパイル済みコードをディスクにキャッシュする

デフォルトでは、Numbaはpythonが再起動するたびに関数を再コンパイルする。cache=Trueを追加してコンパイル済みバイナリをディスクに保存しよう：

@njit(cache=True)
def heavy_computation(arr):
    result = np.zeros_like(arr)
    for i in range(len(arr)):
        result[i] = arr[i] ** 2 + arr[i] * 3 + 1.0
    return result

初回実行：コンパイルして保存。それ以降の実行——Pythonを再起動した後でも——はキャッシュされたバイナリを即座に読み込み、再コンパイルコストはゼロだ。

Numbaが得意なこと（そして苦手なこと）

Numbaが得意なこと：

  • NumPy配列に対するカスタム数値ループ
  • 数学的に重い関数（三角関数、指数、数値の条件分岐）
  • 多くの分岐ロジックを持つシミュレーション
  • きれいにベクトル化できない反復アルゴリズム

Numbaが苦手なこと：

  • あらゆる種類の文字列処理
  • PandasのDataFrame——まず.valuesでNumPy配列を取り出すこと
  • Pythonのリストのリスト——代わりに2次元NumPy配列を使うこと
  • 複雑なメソッドを持つカスタムPythonクラス

ボトルネックがPandasのデータ操作や文字列処理なら、Numbaは適切なツールではない。ループ内の純粋な数値計算では、Pythonエコシステムに並ぶものはない。

パフォーマンス向上の検証とモニタリング

誇大広告を信じる前にベンチマークを取る

誰かの言葉を信じる前に自分で試してみよう。純粋なPython、NumPy、Numbaを同じタスクで直接比較する——あなたの実際のマシンでの実際の数値だ：

import numpy as np
from numba import njit
import time

def pure_python_sum(arr):
    total = 0.0
    for x in arr:
        total += x ** 2
    return total

def numpy_sum(arr):
    return np.sum(arr ** 2)

@njit(cache=True)
def numba_sum(arr):
    total = 0.0
    for i in range(len(arr)):
        total += arr[i] ** 2
    return total

data = np.random.rand(5_000_000)

# Numbaのウォームアップ（初回呼び出しでコンパイル）
numba_sum(data)

for name, fn in [("純粋なPython", pure_python_sum), ("NumPy", numpy_sum), ("Numba", numba_sum)]:
    start = time.perf_counter()
    for _ in range(5):
        fn(data)
    elapsed = (time.perf_counter() - start) / 5
    print(f"{name}: {elapsed:.4f}s")

Python 3.11を実行する一般的な8コアマシンでは、このような結果が期待できる：

純粋なPython: 1.8200s
NumPy:        0.0120s
Numba:        0.0035s

NumPyのarr ** 2はメモリに中間配列を作成するため、ここではNumbaがNumPyを上回る。Numbaは余分なメモリ確保なしに単一パスで二乗和を計算する。

計測前に必ずウォームアップを行う

最初のNumba呼び出しをベンチマークに含めてはいけない。コンパイルは初回呼び出し時に発生する——関数の複雑さによって1〜5秒かかる場合がある。必ずウォームアップ呼び出しを先に実行しよう：

# 小さなスライスでウォームアップ
numba_sum(data[:100])

# 本番の実行を計測
start = time.perf_counter()
result = numba_sum(data)
print(f"経過時間：{time.perf_counter() - start:.4f}s")

cache=Trueを使えば、コンパイルのペナルティは一度だけ発生する——永遠に。その後は、次以降のすべてのPythonセッションでキャッシュされたバイナリが即座に読み込まれる。

Numbaが推論した型を確認する

期待した数値に達していない？関数に対してNumbaが推論した型を確認しよう：

numba_sum.inspect_types()

これはすべての変数の推論された型を出力する。reflected listやobjectが表示されたら？Numbaがその変数でオブジェクトモードにフォールバックしている。NumPy配列または明示的に型付けされたデータに切り替えて修正しよう。

実際のテスト：モンテカルロシミュレーション

モンテカルロによる円周率推定は古典的なストレステストだ。Numbaで1億サンプルを使うとどうなるか：

import numpy as np
from numba import njit, prange
import time

@njit(parallel=True, cache=True)
def monte_carlo_pi(n_samples):
    inside = 0
    for i in prange(n_samples):
        x = np.random.random()
        y = np.random.random()
        if x**2 + y**2 <= 1.0:
            inside += 1
    return 4.0 * inside / n_samples

# ウォームアップ
monte_carlo_pi(1000)

start = time.perf_counter()
pi_estimate = monte_carlo_pi(100_000_000)
elapsed = time.perf_counter() - start

print(f"π ≈ {pi_estimate:.6f}")
print(f"経過時間：{elapsed:.2f}s")

1億サンプルを2秒以内で処理。純粋なPythonの同等のループは3分以上かかる。これが一晩かかるバッチジョブをコーヒーを飲む合間に実行できるほどの差だ。

このパターンはどこでも成立する。配列に対して数学演算をするPythonのループはすべて候補だ。@njit(cache=True)を追加し、入力がNumPy配列であることを確認すれば、純粋なPythonに対して通常50〜200倍の速度向上が見込める。最も遅い関数から始め、ベンチマークを取り、そこから外に向かって作業を進めよう。

Share:
Tags: