Claude CodeでNumba高速化 — プロファイルと結果一致を先に固定する

Pythonの二重三重のforループが遅いとき、Claude Codeに「Numbaで速くして」とだけ頼むと、@njit を付けた時点で終わりになりがちです。速くなったのか、答えが変わっていないのか、どちらも測らないまま差分だけが残ります。

先にプロファイルで遅い関数を特定し、元の出力を基準に残し、@njit 化したあとで np.allclose で一致を確認し、初回コンパイル込みの時間と2回目以降の時間を分けて測る。この順を固定すれば、高速化の成否が数字で判定できます。

本記事は、手元のPython 3.12.2・Numba 0.68.1・NumPy 2.5.4で動かした例を使い、この4段をClaude Codeに任せる指示と検証スクリプトをまとめます。

高速化の前に決めておく4段の順序

Numbaは、型が決まった関数を初回呼び出し時に機械語へコンパイルします。そのため、順序を間違えると「速くなったように見える」か「遅くなったように見える」のどちらかに偏ります。

Claude Code側の作業の組み立ては、公式のワークフロー集にある「バグ修正」「テスト」の型と同じです。再現コマンドを伝える、テストを実行して失敗を直させる、という流れで、高速化では「再現コマンド」が計測スクリプトに、「テスト」が一致確認に置き換わります。

手順1: cProfileで遅い関数を特定する

例として、3次元の点群で各点から最も近い点までの距離を求める関数を使います。三重ループで書いた素直な実装で、1500点を入力にします。

# slow.py(抜粋)
import numpy as np
 
 
def pairwise_min_dist(points):
    n = points.shape[0]
    out = np.empty(n)
    for i in range(n):
        best = np.inf
        for j in range(n):
            if i == j:
                continue
            d = 0.0
            for k in range(points.shape[1]):
                diff = points[i, k] - points[j, k]
                d += diff * diff
            if d < best:
                best = d
        out[i] = best ** 0.5
    return out

プロファイルは次のコマンドで取ります。

python -m cProfile -s tottime slow.py | head -8

手元の出力では、全体約1.9秒のうち pairwise_min_dist の tottime が1.861秒でした。ほかの行はモジュール読み込み由来の0.013秒以下で、遅いのはこの1関数だけです。ここで「どの関数を速くするか」が数字で決まります。

Claude Codeには、この出力をそのまま見せて対象を選ばせます。プロファイルを取らずに関数を選ばせると、見た目が重そうなだけの関数にデコレータを付ける作業になりがちです。

手順2: 元の出力を基準として残す

@njit を付けると、元のPython関数は上書きされるか、別名に隠れます。比較の基準を失わないよう、先に固定シードの入力と出力を残します。

# check.py(抜粋)
rng = np.random.default_rng(0)
pts = rng.random((1500, 3))
ref = slow.pairwise_min_dist(pts)   # デコレータ無しの元の実装

Numba化は、元の関数を書き換えずに別名で作るほうが比較しやすい形です。

# fast.py
from numba import njit
import slow
 
pairwise_min_dist_nb = njit(slow.pairwise_min_dist)

slow.py を触らないので、ref は常に元の実装から出ます。Numba 0.59.0以降は @jit の既定がnopythonモードなので、@njit と @jit の差は気にしなくて構いません。古い記事やコードに @jit(nopython=True) が残っているのは、この変更の前の推奨だったからです。

手順3: 結果一致をnp.allcloseで確認する

一致の確認は、許容誤差を数値で指定して行います。何も指定せず == で比べると、浮動小数点の最後の桁の違いで失敗しますし、逆に緩すぎる許容誤差では変化を見逃します。

python check.py

check.py の比較部分は次の形です。

got = pairwise_min_dist_nb(pts)
print("allclose:", np.allclose(ref, got, rtol=1e-12, atol=0.0))
print("max abs diff:", float(np.max(np.abs(ref - got))))

手元の出力は allclose: True、max abs diff: 6.94e-18 でした。許容誤差は rtol=1e-12 と厳しく取っています。この例は加算の順序を変えない素直な変換なので、ほぼ誤差ゼロになります。

一致しないときの切り分けは、次の順が早い方法です。

• 入力の型が違わないか(float32 と float64、整数配列の混入)。pts.dtype をprintし、元実装とnjit版に同じ配列を渡して確かめます
• 配列を関数内で書き換える実装で、呼び出し間に状態が残っていないか。入力を pts.copy() で複製して渡し、呼び出しごとに元の配列が変わっていないかを見ます
• fastmath や parallel を付けていないか(次節)。デコレータの引数を @njit だけに戻して再実行すれば、原因がここかどうかを切り分けられます

手順4: 初回コンパイル込みの計測を分ける

Numbaの公式ガイドは、性能計測でよくある誤りとして、初回呼び出しのコンパイル時間を実行時間に混ぜることを挙げています。そこで、初回と2回目を別の行で計測します。

t = time.perf_counter(); ref = slow.pairwise_min_dist(pts)
t_py = time.perf_counter() - t
t = time.perf_counter(); got = pairwise_min_dist_nb(pts)
t_first = time.perf_counter() - t
t = time.perf_counter(); got = pairwise_min_dist_nb(pts)
t_second = time.perf_counter() - t

手元の結果は次のとおりです。

初回でもすでに元の約8分の1で、2回目は約400分の1です。ただし、初回の0.222秒のうち実行は数ミリ秒で、大半はコンパイルです。呼び出し回数が少ないスクリプトでは初回の数字が体感を決め、何千回も呼ぶ処理では2回目の数字が効く、という違いが出ます。

複数回呼ぶ形の厳密な比較には、ガイドが勧めるとおり timeit で反復計測するのが確実です。単発のスクリプトで使うなら、@njit(cache=True) でコンパイル結果をディスクに残せます。手元で同じスクリプトを2回続けて動かすと、初回呼び出しが0.185秒から0.096秒に縮みました。

Claude Codeに固定させる指示とCLAUDE.md

手順を毎回口頭で伝えるのは手間なので、CLAUDE.mdに規約として書いておきます。次は例示の断片で、プロジェクトに合わせて許容誤差などを変えます。

## Numbaで数値処理を速くするとき
- 変更前に `python -m cProfile -s tottime <script>` を実行し、自己時間が
  最大の関数だけを対象にする
- 元の関数は書き換えず、`njit(関数)` を別名で作って比較する
- 固定シードの入力で `np.allclose(ref, got, rtol=1e-12, atol=0.0)` を実行し、
  False なら速度の話をせず原因を先に調べる
- 計測は初回(コンパイル込み)と2回目以降を別々に報告する
- `fastmath=True` と `parallel=True` は、一致確認を通した後に1つずつ足す

依頼の文面は、公式のテスト作業の型に沿って「実行して失敗を直す」まで含めると、往復が減ります。

slow.py の pairwise_min_dist を njit 化して。
手順は CLAUDE.md の規約どおり。check.py を実行して、
allclose が True になり、初回と2回目の時間が出るところまで確認して。

応答の文面は実行ごとに変わります。判定に使うのは check.py の出力で、allclose がFalseなら作業が先へ進まない構成にしておきます。

fastmathとparallelは一致確認のあとに足す

Numbaには fastmath=True と parallel=True があり、公式の性能ガイドでは総和の例で fastmath により35.2msが17.8msに、parallel=True と prange の組み合わせで9.81msに縮んでいます。数値は公式がi7-4790で測ったもので、手元の環境の結果とは別物です。

ただし、fastmath は浮動小数点の再結合を許す設定で、結果が変わり得ます。ガイドには、(x - y) + y を x=0、y=inf で計算すると、fastmath=False なら nan、True なら 0.0 になる例が載っています。一致確認の前に付けると、遅さと誤差の原因が混ざります。

一致確認を通した関数に、オプションを1つずつ足して check.py を回し直す。この運用なら、どのオプションで答えが動いたかが分かります。parallel=True は総和のような縮約で加算順が変わるため、allclose の許容誤差を見直す対象にもなります。

つまずきやすい点

• Numbaが対応しない書き方: pandasのDataFrameや任意のPythonオブジェクトを渡す関数は、nopythonモードでコンパイルできません。ガイドは、対応していないコードをobjectモードに落とす方法は最適な速度につながらないと説明しています。NumPy配列とスカラーを受け取る関数に切り出してから適用します。
• キャッシュの古さ: cache=True は、別ファイルからimportした関数の変更を検出できません。その場合は古いコードが使われ得るため、計測の前にキャッシュを消します。キャッシュの置き場は、既定ではソースファイルと同じ場所の __pycache__ ディレクトリです。NUMBA_CACHE_DIR で場所を変えている場合は、そちらを消します。グローバル変数も定数としてキャッシュされます。
• デバッグしづらさ: 環境変数 NUMBA_DISABLE_JIT=1 を付けると、デコレータが何もしなくなり、元のPython関数が呼ばれます。値を print したり pdb で止めたりしたいときは、この設定で動かします。

NUMBA_DISABLE_JIT=1 python check.py

向いている処理と、先に別の手を打つ処理

Numbaは、NumPy配列とループが中心のコードで効きます。逆に、すでにNumPyのベクトル演算で書かれている関数に @njit を付けても、速くなるとは限りません。公式ガイドの例では、ベクトル化した版は @njit の有無で0.581秒と0.659秒とほぼ同じで、forループ版だけが25.2秒から0.670秒へ縮んでいます。

プロファイルの結果、遅い関数がforループではなくI/Oや外部ライブラリの待ちなら、Numbaの出番ではありません。PyTorchならDataLoaderの待ちを直す手順、Node.jsならcpuprofileから遅い関数を絞る手順が別の入口になります。

なお、MATLABのコードをPythonに移した直後で、まず結果の一致を取りたい場合はMATLABからPythonへの移植と検算の流れが先です。移植で数値が合ってから、本記事の手順でループを速くすると、原因の切り分けが楽になります。

まとめ

Numba化の成否は、デコレータを付けたかではなく、allclose が通ったあとの2回目の時間で決まります。プロファイル、基準出力、一致確認、初回と2回目の分離計測を、CLAUDE.mdとチェックスクリプトに寄せておけば、Claude Codeに任せる範囲が広がっても、変わってはいけない出力の側が先に壊れたと気づけます。