Numbaを用いたPython/Numpyの高速化

f:id:k-kty:20190216135709p:plain

まえがき

機械学習関連のプロジェクトではPythonを使う人が多いと思うのですが、Pythonでの処理って遅いですよね。

Tensorflowでの学習などにおいてはPython側で計算をしているわけではないので大丈夫なのですが、前処理の際などにPythonの実行速度の遅さがとても気になります。コンパイラ言語とスクリプト言語を比べるのは酷ですが、「C++なら数秒で終わるのになあ」という処理も数分かかってしまったりします。

もちろんC++のプログラムに匹敵するほどの速度を出すのは無理ですが、できるだけ高速化したいですよね。そのための方法はいろいろあるのですが、今回はNumbaというJITコンパイラを用いる方法を紹介したいと思います。

Numbaとは

公式サイトによると、

Numba is an open source JIT compiler that translates a subset of Python and NumPy code into fast machine code.

「Python/Numpyコード（の一部）を、高速な機械語に変換するためのオープンソースJITコンパイラ」ってことみたいですね。

もう少し細かい特徴としては以下のような感じです。

関数にデコレータをつけるだけで動作する
Numpyの配列や関数を用いた処理に特化している
ランタイム時に、最適化された機械語を作成して実行してくれる
Numbaを利用するためにPythonのインタープリターを変えたり、コードの実行前にコンパイルするようにしたりする必要がない
CPU/GPUを用いた並列計算のための機能もある

実際に使ってみる

高速化前のコード

以下のような関数があったとします。配列（arrとします）を受け取り、[log(arr[1]/arr[0]), log(arr[2]/arr[1])... ] のような（長さがarrより1短い）配列を返します。

def calculate_diffs(arr):
  ret = numpy.empty(arr.size - 1)
  for i in range(arr.size - 1):
    ret[i] = numpy.log(arr[i + 1] / arr[i])
  return ret

以下のコードを用いて測定をすると（かなり雑ですが、今回はだいたいのパフォーマンスの比較ということで...）、出力は 183.20000410079956 となりました。約3分です。なお関数内のnumpy.empty(...)の部分は計算とは関係ありませんがe-5秒くらいのオーダーだったので無視してよさそうです。（ちなみにですがnumpy.random.rand(...)の部分は無視できないほど時間がかかります。）

import time
import numpy

# 以上の関数定義

if __name__ == '__main__':
  arr = numpy.random.rand(10 ** 8)

  start_t = time.time()
  calculate_diffs(arr)
  end_t = time.time()

  print(end_t - start_t)

JITを有効化してみる

以下のようにデコレータをつけて、JITを有効化してみます。

@numba.jit(nopython=True)
def calculate_diffs(arr):
  ret = numpy.empty(arr.size - 1)
  for i in range(arr.size - 1):
    ret[i] = numpy.log(arr[i + 1] / arr[i])
  return ret

測定の結果、出力は1.805896282196045となり、100倍ほどの高速化に成功しました。

JITを有効化/並列化してみる

以下のようにデコレータを変更し、range関数をnumba.prange関数に置き換え、並列化してみます。 ret = numpy.empty(arr.size - 1)の部分は、numbaが勝手にうまいこと(hoist)してくれて、期待通りの動作をします。

@numba.jit(nopython=True, parallel=True)
def calculate_diffs(arr):
  ret = numpy.empty(arr.size - 1)
  for i in numba.prange(arr.size - 1):
    ret[i] = numpy.log(arr[i + 1] / arr[i])
  return ret

測定の結果、出力は0.7335679531097412となり、200倍以上の高速化に成功したことになります。

まとめ

限られたケースにはなりますが、Numbaをうまく使えば非常に簡単に高速化ができることがわかりました。