numbaというライブラリを使うと、Pythonのコードを比較的簡単に高速化できます。
うまくいけば、from numba import jitを書いて、高速化したい関数の前の行に@jitを書くだけで高速化できます。
仕組みとしては、numbaはPythonの仮想マシンコードを取得し、LLVM IRにコンパイルし、LLVMを使ってネイティブコードにするようです。
初回実行時は、コンパイル処理が走るので、若干遅くなりますが、重い処理だと、コンパイル時間を考えてもnumbaの方が速いこともあります。
利点と欠点
先に述べておきます。
利点
- 場合によっては、コード自体は改造せずに手軽に高速化できる
- コードの改造があったとしても、軽微な改造で済むことも多い
- 別ファイルに分けてビルドする、みたいな手間なことが必要なく、手軽に
.pyファイルの中で使える
欠点
- すべてのPython機能がサポートされているわけではないらしく、場合によっては
@jitをつけるだけでは済まない- 型の扱いが厳格になるので、今まで「なぁなぁ」で動いていた部分がエラーになったりするようになります。また、numbaにうまく型推論してもらうよう工夫が必要になる場合があります。
- 当然だが、動かすのにnumbaが必要になる。環境によっては、numbaのインストールに苦労する
- conda環境だと簡単なようです。pipで入る環境も結構あります。私の使っているArch Linuxでは、今現在、Python 3.8、LLVM 9.0になっており、どちらも現時点ではnumbaが対応していないため、ビルドするのは諦めてDockerでconda環境を使ってます
- コンパイルに時間がかかるので、闇雲につけたら逆効果になる
例
とてもうまくいく事例を紹介します。
なんかしらんけど、とっても遅い関数があります。
import sys
sys.setrecursionlimit(100000)
def ack(m, n):
if m == 0:
return n + 1
if n == 0:
return ack(m - 1, 1)
return ack(m - 1, ack(m, n - 1))
ちょっと時間測ってみますね。
import time
from contextlib import contextmanager
@contextmanager
def timer():
t = time.perf_counter()
yield None
print('Elapsed:', time.perf_counter() - t)
with timer():
print(ack(3, 10))
8189
Elapsed: 10.270420542001375
10秒もかかっちゃいました。
数字を大きくするともっと時間かかりますが、本当に時間がかかるのでおすすめしません。
特に、3を4に増やすと、多分、死ぬまでかかっても終わらないので、全くおすすめしません。
この関数はアッカーマン関数として知られています。
これをnumbaで速くしてみます。
from numba import jit
@jit
def ack(m, n):
if m == 0:
return n + 1
if n == 0:
return ack(m - 1, 1)
return ack(m - 1, ack(m, n - 1))
# 1回目
with timer():
print(ack(3, 10))
# 2回目
with timer():
print(ack(3, 10))
# 3回目
with timer():
print(ack(3, 10))
8189
Elapsed: 0.7036043469997821
8189
Elapsed: 0.4371343919992796
8189
Elapsed: 0.4372558859977289
なんと、10秒かかっていたものが、初回0.7秒、2回目以降0.4秒にまで縮みました。
1行付け足すだけでこうなるなら、本当にお得ですね。
思ったより速くない場合は
Objectモードが使われている可能性があります。
numbaには、No PythonモードとObjectモードがあり、一旦No Pythonモードでコンパイルされますが、失敗したらObjectモードでコンパイルされます。
(ただし、この仕様は将来なくなり、デフォルトでNo Pythonモードのみ、Objectモードはオプションになるようです)
前者は、型が全部直接扱われるのに対し、後者は、PythonのオブジェクトをPython C APIから叩き、前者のほうがより高速です。
さらに、後者だと、ループを効率よくネイティブコードに書き換えられない場合がありますが、前者だとループが効率化されます。(loop-jitting)
No Pythonモードを強制するためには@jit(nopython=True)または@njitと書きますが、こうすると「型が分からない」といった類のエラーがよく発生するようになります。
基本的には、
- ひとつの変数には、ひとつの型のみを入れる
-
returnが複数ある場合、どのreturnでも同じ型を返すようにする - No Pythonモードの関数で呼び出している関数もNo Pythonモードにする
などの方法で、型を明確にすることができます。
比較的簡単にNo Pythonモードに対応させるには
- 計算が重い部分だけを別関数に切り出してnumbaに対応させる
- 切り出した中でも、型をごちゃごちゃいじったり、オブジェクトにあれこれ代入したりしている部分は、外に出すことを考える
- numba対応する部分の前処理あるいは後処理の形で、なんとかnumbaでやらずに済ませることを考える
などがおすすめです。
Pythonが得意なことはPythonでやる、Pythonが苦手なことはnumbaでやる、を意識します。
並列化してみる
@jit(parallel=True)とすると、forループでrangeの代わりにprangeが使えます(from numba import prangeが必要)。
prangeで書いたループは、並列化されます。
コンパイル結果のキャッシュ
@jit(cache=True)で、コンパイル結果をキャッシュファイルに書き出し、毎回コンパイルする手間を避けることができます。
fastmathを使う
@jit(fastmath=True)で使えます。
gccやclangにもあるfastmathを有効にします。やや危なっかしい最適化でfloatの計算を早くするやつです。
CUDAを使う
あまりお手軽ではありませんが、一応、CUDAも使えます。
CUDA使ったことある人なら、以下のコードでなんとなく分かるかと思います。
個人的には、これだったらcupyでいいかなぁ、と思いました。
import numpy as np
from numba import cuda
@cuda.jit
def add(a, b, n):
idx = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
if idx < n:
a[idx] += b[idx]
N = 1000000
a_host = np.array(np.ones(N))
b_host = np.array(np.ones(N))
a_dev = cuda.to_device(a_host)
b_dev = cuda.to_device(b_host)
n_thread = 64
n_block = N // n_thread + 1
add[n_block, n_thread](a_dev, b_dev, N)
a_dev.copy_to_host(a_host)
print(a_host) # Expect: [2, 2, ..., 2]
もしlibNVVMがない、などと怒られたら、CUDAをインストールしていない(conda install cudatoolkitなどでインストールできます)か、環境変数の設定が必要です。
Google Colabなどでの設定例:
import os
os.environ['NUMBAPRO_LIBDEVICE'] = "/usr/local/cuda-10.0/nvvm/libdevice"
os.environ['NUMBAPRO_NVVM'] = "/usr/local/cuda-10.0/nvvm/lib64/libnvvm.so"
まとめ
numbaを使うと、ざっくりとPythonコードを高速化できます。
簡単な使い方を見ていきました。
参考になれば幸いです。
参考文献
公式ドキュメントは、必要なところだけ読めば、さほど長くなく、非常に参考になります。
特にPerformance Tipsは役立ちます。
CUDAを使う際のの環境変数設定については、
https://colab.research.google.com/github/cbernet/maldives/blob/master/numba/numba_cuda.ipynb
を参考にしました。