NumPyをつかおう

  •  
 
abe2018年12月5日 - 20:28 に投稿

タグ

発端

トビウオさんの記事を興味深く読んだのだが、最後の「例:録音した音声の音量(RMS)を計算して随時表示する」で、2点ほど気になったのである:

  1. RIFFヘッダーならともかく、ただのPCMデータ部分をstructで処理するのは大げさではないか?
  2. dB値計算はリアルタイムであるべきであり、pure-Pythonで書いては処理が追いつかなかったりするのではないか?

私は普通のさらりーぱーそんでありながら、なぜかNumPyを触る仕事が多く、こういうプログラムはどうしてもNumPy化したくてしかたがない。

という訳で、structをやめることでどれくらいシンプルに書けるのか、また、処理速度がどこまで速くなるのかを検証してみよう。

準備&トビウオさん実装の検証

まずは準備。 私がいじったことで、トビウオさん実装と違う動作をすることがないかの回帰試験と、トビウオさん実装よりもどれだけ速く(あるいはひょっとすると遅く)なったのかを確認できるような仕掛けを用意しておこう:

#!/usr/bin/env python

import math
import struct
import timeit

import numpy as np


NAME2IMPL = {}
def impl(name):
    def decorator(method):
        global NAME2IMPL
        NAME2IMPL[name] = method
    return decorator


@impl('original')
def original(in_data):
    # bytes型を配列に変換する
    # (とりあえず8bit・モノクロだとした例を書く。
    # データは1バイトづつであり、0〜255までで中央値が128であることに注意)
    in_data2 = struct.unpack(f'<{len(in_data)}B', in_data)
    in_data3 = tuple((x - 128) / 128.0 for x in in_data2)

    # 読み取った配列(各要素は-1以上1以下の実数)について、RMSを計算する
    rms = math.sqrt(sum([x * x for x in in_data3]) / len(in_data3))

    # RMSからデシベルを計算して表示する
    db = 20 * math.log10(rms) if rms > 0.0 else -math.inf

    return db


def test():
    original = NAME2IMPL['original']

    size = 4096
    test_patterns = [
        np.full(size, 255, dtype=np.uint8).tobytes(),
        np.full(size, 0, dtype=np.uint8).tobytes(),
        np.full(size, 128, dtype=np.uint8).tobytes(), ]
    test_patterns += [
        np.random.RandomState(seed).randint(0, 255, size=size, dtype=np.uint8).tobytes()
        for seed in range(10) ]
    for name, func in NAME2IMPL.items():
        print(f'testing {name}', end='', flush=True)
        for i, test_pattern in enumerate(test_patterns):
            expected = original(test_pattern)
            actual = func(test_pattern)
            assert np.allclose(actual, expected), f'{name} yields incorrect result for pattern #{i}'
            print('.', end='', flush=True)
        print('ok')


def benchmark():
    original = NAME2IMPL['original']

    size = 4096
    num_trials = 20000
    in_data = np.random.uniform(0, 255, size=(2**12,)).astype(np.uint8).tobytes()
    for name, func in NAME2IMPL.items():
        #if name == 'original': continue
        func(in_data)
        t = timeit.timeit(lambda: func(in_data), number=num_trials)
        print(f'{name}\t{1e6*t/num_trials} us/iter.')

def main():
    test()
    benchmark()

if __name__ == '__main__':
    main()

# the end of file

そんなに難しいコードではなかろう。

test()では、いくつかのパターン(「0ばかり」「128ばかり」「255ばかり」は境界ケース、残りは一定シードの乱数)で、トビウオさん実装(original())と同じ結果になることを確認している。

また、benchmark()では、各実装を4 KiBの乱数列を引数として20,000回呼んで平均処理時間(μs単位)を計算している。 本当ならば標準偏差ないし標準誤差くらい計算しておきたいものだが、面倒なので省略する。

手許のマシン(Intel® Xeon® E3-1275v6)で実行すると、こんな感じ:

実装名 処理時間(@4 KiB; μs)
original() 457.6

仮に元の音源が48 KHzとすると4,096サンプルは85 ms (=85,000 μs)程度なので、pure-Pythonであっても全く性能問題はないことが判る。 疑ってすみません > トビウオさん

NumPy化する

記事のタイトルに反して、NumPyのなんたるかについては説明しない。

仮に知らなくとも、取り敢えず「数値計算をお手軽かつ高速にするためのライブラリーだよ」という程度の認識でよろしい。 特に、今回のように「同じ型の数値がたくさん並んでいる」というケースを、ひょっとすると恐ろしいくらい高速に実行できるライブラリーである。

NumPyを使ってトビウオさん実装を素直に書き換えるとこんな感じになる:

@impl('NumPy')
def by_numpy(in_data):
    # in_dataを、符号なし8ビット整数(np.uint8)の配列とみなす
    in_uint8 = np.frombuffer(in_data, dtype=np.uint8)
    # 128を引き(-128.0)、更に128で割る(/128.0)ことで、(だいたい) [-1.0,1.0]の範囲となるように正規化する
    in_normalized = (in_uint8-128.0)/128.0
    # 各要素を自乗(**2)し、その平均を計算(.mean())し、ついでに平方根をとる(np.sqrt())
    rms = np.sqrt((in_normalized**2).mean())
    # エネルギー比をdBに換算(20*np.log10())する
    return 20 * np.log10(rms)

ご覧の通り、リスト内包表記がきれいに消えた。NumPyを使い慣れているせいか、私にとってはずいぶん読みやすくなったように感じられる。

これを実行すると:

実装名 処理時間(@4 KiB; μs)
numpy() 18.86
original() 457.6

なんとこれだけで24倍速。

これだけ高速化した理由はいくつかあるけど、代表的なのは以下のような感じ:

  1. 評価されるPython式の個数が減った: トビウオさん実装では、リスト内包表記内で__sub__()だの__truediv__()だのが何度も評価され、そのたびにPythonインタープリターが頑張っていた。NumPyだと__sub__()__truediv__()自体は1回だけで、各要素の計算はCコード内で行われるので、Pythonインタープリターはほとんど何もしないで済む。
  2. その引き算や割り算自体が高速: トビウオさん実装では、一要素ごとに__sub__()などが行われていたが、NumPyだとCコード内で行われており、それはSIMD命令を用いて複数要素まとめて効率的に実行されている可能性が高い
  3. メモリーコピーが減り、キャッシュヒット率が高い: トビウオさん実装では元のバイト列を翻訳・コピーしたような配列を作っていたが、numpy.frombuffer()はコピーせずにラップするので4 KiBのメモリーコピーが発生しない

NumPyのまま、更に高速化

単純なNumPy化には、算数的・コンピューター的に考えていろいろと無駄が多い:

  1. 「割ってから平均」とすると配列の全要素を割る必要があるが、「平均をとってから割る」とすれば割り算は1回で済む
  2. 「平方根をとってからlog10()」だと平方根演算が必要だが、「log10()を取ってから0.5倍」とすれば(より簡単な演算である)掛け算にできる
  3. 「全要素を自乗した配列を作ってから平均」とすると中間配列を作る必要があるが、「配列自身の内積」とすれば中間配列は不要

最初の1つは小学生並だが、残りの2つは高校生の数学が必要。とはいえ、そんなに難しい話ではあるまい。

これを実装するとこんな感じ:

@impl('NumPy (fast)')
def by_numpy_fast(in_data):
    dtype = np.float64
    in_uint8 = np.frombuffer(in_data, dtype=np.uint8)
    in_0centered = np.subtract(in_uint8, 128, dtype=dtype)
    mean_square = np.dot(in_0centered, in_0centered) / (in_0centered.size*128*128)
    return 10 * np.log10(mean_square)

プログラムが多少読みづらくなってしまった。が、実行してみるとそうしただけの効果があったことが判る:

実装名 処理時間(@4 KiB; μs)
numpy_fast() 9.268
numpy() 18.86
original() 457.6

numpy()から倍速化した。 これは主に、先の3つ目に書いた「中間配列は不要」というのが効いていると思われる。 数値計算では、計算回数よりも先に、メモリーアクセスがネックになることが多いのだ。

Numbaを使う

このようにプログラムを高速化していると、そのうち「えーい。C言語(やインラインアセンブリー)で直接書いた方が速いではないか!」という気分になってくる。 今回の場合も、NumPyにとどまる限りはin_0centeredという中間配列を作らざるを得ず、どうにかCで書きたいなぁ…と思えてくる。

Pythonには、ネイティブコードを使う仕組みが複数用意されている:

  1. Python自体のC拡張: Python本体に備わっている、C言語などでPythonの拡張モジュール(*.so, *.dylib, *.dll)を作る仕組み。目的のロジック自体はともかく、Pythonインタープリターとおしゃべりする部分を書くのがチョー面倒
  2. Cython: その「Pythonインタープリターとおしゃべりする部分」を、Pythonっぽい文法の言語で書けるようにする。楽にはなるけど、やはり*.soなどを作らなくちゃならないので、ちょっと面倒
  3. Numba: Python関数にアノテーション(@numba.jit)をつけるだけで、動的にネイティブコードコンパイルしてくれる。とっても楽だけど、「あのSIMD命令を使いたい」みたいな制御はできない

拡張モジュールを作るような時間的・精神的な余裕はないので、最後のNumbaを使ってみよう:

import numba

@numba.jit
def _by_numba_aux(in_uint8):
    dtype = np.float64
    sum_square = dtype(0)
    for i in range(in_uint8.size):
        tmp = dtype(in_uint8[i]) - dtype(128)
        sum_square += tmp*tmp
    mean_square = sum_square / dtype(in_uint8.size*128*128)
    return dtype(10) * np.log10(mean_square)

@impl('Numba')
def by_numba(in_data):
    return _by_numba_aux(np.frombuffer(in_data, dtype=np.uint8))

やったことは2つ:

  1. 高速化したいPython関数に、@numba.jitというアノテーションを付けた
  2. NumPyのnumpy.dot()などは使わないべたべたのforループとし、中間配列を一切なくした

トビウオさん実装→NumPy化したとき、高速化の要因に「評価されるPython式の個数が減った」というのを挙げた。 その観点からすると、上記2点目の「べたべたのforループ」というのは許しがたい悪手に見える。 しかしNumbaは、そのようなforループすら綺麗に最適化してくれる(というか、むしろforループ化してNumbaに最適化したほうが効率的になる)のである:

実装名 処理時間(@4 KiB; μs)
numba() 5.026
numpy_fast() 9.268
numpy() 18.86
original() 457.6

numpy_fast()から更に倍速化。 これはもちろん、in_0centeredという中間配列を作らなくて済むようにしたのが理由だろう。

Numba版を更に高速化

さきほどはNumbaをデフォルトオプションのまま利用したが、実はコンパイルオプションを選べるようになっている。 今回のように「Python関数内では、NumPy以外のPythonオブジェクトを扱わない」「排他制御(GIL; Global Interpreter Lock)は不要である」の場合、その旨をアノテーションに記述するだけでより高速に実行できる(かもしれない)ネイティブコードが生成されることが期待できる:

@numba.jit(fastmath=True, nogil=True, nopython=True)
def _by_numba_opt_aux(in_uint8):
    dtype = np.float64
    sum_square = dtype(0)
    for i in range(in_uint8.size):
        tmp = dtype(in_uint8[i]) - dtype(128)
        sum_square += tmp*tmp
    mean_square = sum_square / dtype(in_uint8.size*128*128)
    return dtype(10) * np.log10(mean_square)

@impl('Numba (optimized)')
def by_numba_opt(in_data):
    return _by_numba_opt_aux(np.frombuffer(in_data, dtype=np.uint8))

今回は「関数内でPythonオブジェクトを使わない(nopython=True)」「排他制御は不要(nogil=True)」「数値計算は精度よりも速度優先で(fastmath=True)」を有効化した。 データ量がもっと多かったり、処理が複雑だったりする場合は「自動的にマルチスレッド化する(parallel=True)」なんかをつけても良いかもしれない。

これを実行すると:

実装名 処理時間(@4 KiB; μs)
numba_opt() 1.452
numba() 5.026
numpy_fast() 9.268
numpy() 18.86
original() 457.6

numba()から更に3倍以上高速化した。 もともとのトビウオさん実装からすると、300倍以上高速化したことになる。

Numbaが出力したコードを確認したが、「Pythonから呼び出す」ということを考えるとほぼ完璧なコードになっており、これ以上の高速化は期待しにくい。 (log10()の計算を自前で近似する、というのが考えられるが、さすがにそこまでいくとNumbaでできるかどうか怪しくなってくる)

結論

当初の目的はこんなのだった:

という訳で、structをやめることでどれくらいシンプルに書けるのか、また、処理速度がどこまで速くなるのかを検証してみよう。

「シンプルに書ける」というのをすっかり忘れていたよ。

これらに関してまとめると、こんなところだろうか:

  1. 単純にNumPy化した場合、(少なくとも慣れている目から見ると)読みやすくなるが、Numbaが入ってくるとだんだん読みづらくなってきた
  2. dB値計算くらいであればそもそも処理速度に問題はなかったが、単にNumPy化するだけで20倍以上速くなるし、本気を出せば(拡張モジュールを作らなくとも) 300倍以上速くなる

あんまりうまい結論じゃなくてすみません。


追記(2018年12月6日): まだまだ高速化するようだ

Numba高速版から、関数呼び出しを除去

これまでは、入り口のby_numba_opt()と実装の_by_numba_opt_aux()とを別関数にしていた。これは、入り口にあるnp.frombuffer()@numba.jit(..., nopython=True, ...)と干渉するのではないかと考えたから。 実際に動かしてみると、np.frombuffer()@numba.jit(..., nopython=True, ...)内に入れても問題なくコンパイルできるらしい:

@impl('Numba (optimized2)')
@numba.jit(fastmath=True, nogil=True, nopython=True)
def by_numba_opt2(in_data):
    dtype = np.float64
    sum_square = dtype(0)
    in_uint8 = np.frombuffer(in_data, dtype=np.uint8)
    for i in range(in_uint8.size):
        tmp = dtype(in_uint8[i]) - dtype(128)
        sum_square += tmp*tmp
    mean_square = sum_square / dtype(in_uint8.size*128*128)
    return dtype(10) * np.log10(mean_square)

実行すると60%ちょっと高速化しており、Pythonの関数呼び出しが、数値演算(積和)よりも3桁程度遅いことが判る:

実装名 処理時間(@4 KiB; μs)
numba_opt2() 0.8862
numba_opt() 1.452
numba() 5.026
numpy_fast() 9.268
numpy() 18.86
original() 457.6

更に、もう少し算数

dB計算には底を10とする対数(常用対数)が必要だが、数学的にはeを底とする対数(自然対数)の方がTaylor展開が綺麗であり、数値計算的にもやりやすそうだ。ちゃんと確認はしていないが、log10(x)は内部でlog(x)を計算した上で底の変換をしているのだろう。これを利用すると、掛け算を一つ減らすことができる:

@impl('Numba (optimized3)')
@numba.jit(fastmath=True, nogil=True, nopython=True)
def by_numba_opt3(in_data):
    dtype = np.float64
    in_uint8 = np.frombuffer(in_data, dtype=np.uint8)
    sum_square = dtype(0)
    for i in range(in_uint8.size):
        tmp = dtype(in_uint8[i]) - dtype(128)
        sum_square += tmp*tmp
    mean_square = sum_square / dtype(in_uint8.size*128*128)
    # dtype(10*np.log10(np.e))の部分はコンパイル段階で定数畳み込みされるだろう
    # →np.log10()内で行われているであろう底変換を、外に括り出して減らすことができるはず?
    return dtype(10*np.log10(np.e)) * np.log(mean_square)
実装名 処理時間(@4 KiB; μs)
numba_opt3() 0.8590
numba_opt2() 0.8862
numba_opt() 1.452
numba() 5.026
numpy_fast() 9.268
numpy() 18.86
original() 457.6

たったの3%か。何度も実行すると時々逆転することもあるのを考えると、確実に効くようなものではないのだろう。


追記(2018年12月7日): 下記2つのtypoを修正した:

  1. トビウオさん実装部分で、def original(in_data):の行を書き忘れていた(コピペミス)
  2. Numba実装部分で、import numbaの行を漏らしていた

NumPyにしただけで24倍速になったのは驚きですね。 より凝った処理をしようとすると、読み込み部分や計算処理はなるべく速くしたいものですので、 こうした記事は勉強になります。

コメントを追加

プレーンテキスト

  • HTMLタグは利用できません。
  • 行と段落は自動的に折り返されます。
  • ウェブページのアドレスとメールアドレスは自動的にリンクに変換されます。
CAPTCHA
この質問はあなたが人間の訪問者であるかどうかをテストし、自動化されたスパム送信を防ぐためのものです。