読者です 読者をやめる 読者になる 読者になる

ゆとりデータサイエンティストの諸々所感

データ分析会社で研究開発をしている、ゆとり世代データサイエンティストが学んだ内容や最新トピックについて諸々語る予定

Python高速化 Numba入門 その5

Numba Python

今回はGPUの数値実験についての記事です。

計算環境

OS: Windows 7 64bit
CPU: Intel Core i7-4771 CPU @ 3.50GHz
GPU: NVIDIA GeForce GTX 650
MEM: 32GB

我が家の計算マシンです。
これをWindowsで構成しているのは、もったいない気がしていますが、
LinuxGPU載せても用途が限られすぎるので妥協です。

環境構築

ここまで、特に環境構築について触れて来ませんでしたがGPUを使うと苦労したので少し書きます。

結論としては、『Anacondaを入れましょう!』が最善手のようです。

Download Anaconda Python Distribution

Win機64bitで環境を揃えるのはかなりめんどくさいです。というか頑張ったんですが
エラーが直らず断念しました

GPU使わないなら、GohlkeさんのページからNumbaとllvmliteを入れればお手軽に動かすことが出来ます。

Python Extension Packages for Windows - Christoph Gohlke

Win機のGPGPUは大体面倒ですが、世のGPUはゲーム用途でかなりWin機に搭載されているので、このミスマッチはいかんともしがたい感じです。

その他ライブラリはこんな感じ。記事を書いている間にnumbaの0.16が出たのでアップデートしています。

  • python=2.7.8
  • numba=0.16.0
  • numpy=1.9.1
  • llvmlite=0.2.0

実際に動かしてみた

その3まで同様に1,000地点間の3次元の距離行列を求めるコードで実験してみます。

# coding: utf-8
import numpy
from numba import double
from numba.decorators import jit
from numba import guvectorize
from numba import cuda
import math
import time

def pairwise_python(X, D):
    M = X.shape[0]
    N = X.shape[1]
    for i in range(M):
        for j in range(M):
            d = 0.0
            for k in range(N):
                tmp = X[i, k] - X[j, k]
                d += tmp * tmp
            D[i, j] = numpy.sqrt(d)
    return D

@jit
def pairwise_numba(X, D):
    M = X.shape[0]
    N = X.shape[1]
    for i in range(M):
        for j in range(M):
            d = 0.0
            for k in range(N):
                tmp = X[i, k] - X[j, k]
                d += tmp * tmp
            D[i, j] = numpy.sqrt(d)


@jit('f8[:, :](f8[:, :], f8[:, :])', nopython=True)
def pairwise_numba_jit(X, D):
    M = X.shape[0]
    N = X.shape[1]
    for i in range(M):
        for j in range(M):
            d = 0.0
            for k in range(N):
                tmp = X[i, k] - X[j, k]
                d += tmp * tmp
            D[i, j] = numpy.sqrt(d)

    return D


@guvectorize(['void(f8[:, :], f8[:, :])'], '(x, y)->(x, x)')
def pairwise_numba_guvec(X, D):
    M = X.shape[0]
    N = X.shape[1]
    for i in range(M):
        for j in range(M):
            d = 0.0
            for k in range(N):
                tmp = X[i, k] - X[j, k]
                d += tmp * tmp
            D[i, j] = numpy.sqrt(d)

@cuda.jit('void(f8[:, :], f8[:, :])')
def pairwise_numba_cuda1(X, D):
    M = X.shape[0]
    N = X.shape[1]
    
    i, j = cuda.grid(2)
    
    if i < M and j < M:
        d = 0.0
        for k in range(N):
            tmp = X[i, k] - X[j, k]
            d += tmp * tmp
            
        D[i, j] = math.sqrt(d)

if __name__ == '__main__':
    N = 1000
    dim = 3
    
    t = time.time()
    numpy.random.seed(0)
    X = numpy.random.random((N, dim))
    D = numpy.empty((N, N))
    pairwise_python(X, D)
    print "python           :", time.time() - t

    t = time.time()
    numpy.random.seed(0)
    X = numpy.random.random((N, dim))
    D = numpy.empty((N, N))
    pairwise_numba(X, D)
    print "numba_jit        :", time.time() - t

    t = time.time()
    numpy.random.seed(0)
    X = numpy.random.random((N, dim))
    D = numpy.empty((N, N))
    pairwise_numba_jit(X, D)
    print "numba_jit_type   :", time.time() - t

    t = time.time()
    numpy.random.seed(0)
    X = numpy.random.random((N, dim))
    D = numpy.empty((N, N))
    pairwise_numba_guvec(X)
    print "numba_guvectorize:", time.time() - t

    
    t = time.time()
    numpy.random.seed(0)
    X = numpy.random.random((N, dim))
    D = numpy.empty((N, N))
    griddim = 100, 100
    blockdim = 16, 16
    pairwise_numba_cuda1[griddim, blockdim](X, D)
    print "numba_cuda       :", time.time() - t

結果は以下、

blockdim: (16, 16)
griddim: (63, 63)
python           : 2.21000003815     #普通のPython
numba_jit        : 0.0870001316071   #型指定なしCPU
numba_jit_type   : 0.0090000629425   #型指定ありCPU
numba_guvectorize: 0.0090000629425   #numpyのGUFunc
numba_cuda       : 0.0139999389648   #GPU

あれ?GPU速くない・・・
問題サイズが小さいからかもとのことで、1,000地点間の距離から10,000地点間の距離に変更して実行してみる。
普通のPythonはまともに計算出来ないのでここで退場

blockdim: (16, 16)
griddim: (626, 626)
numba_jit        : 0.734000205994  #型指定なしCPU
numba_jit_type   : 0.90399980545   #型指定ありCPU
numba_guvectorize: 0.888999938965  #numpyのGUFunc
numba_cuda       : 0.59299993515   #GPU

3割程度速くなった。GPUを導入したからといって劇的な改善があるわけではなさそう。

また、意図的にブロック数を増やしてみると

blockdim: (16, 16)
griddim: (5000, 5000)
numba_jit        : 0.733999967575
numba_jit_type   : 0.888999938965
numba_guvectorize: 0.858000040054
numba_cuda       : 2.18400001526

オーバーヘッドでかなり遅くなることがわかった。GPU君は高速化には寄与するものの、しっかりしたチューニングが必要そうな感じです。

Shared Memoryの効果検証

前回も触れた、行列をブロックに分けてそれをShared Memoryに格納して行列積を行う方法を検証してみる。
しかし、こいつはブロックの分けるのに端数が出ると、実装が面倒なので、

今回はブロックサイズ16、ブロック数50として行列サイズN=50×16=800の正方行列の積で検証する

実行したコードが以下

# coding: utf-8
import numpy
from numba import double
from numba.decorators import jit
from numba import guvectorize
from numba import cuda
import math
import time
from numba import float32, float64

bpg = 50
tpb = 16
n = bpg * tpb

#@cuda.jit('void(f4[:, :], f4[:, :], f4[:, :])')
@cuda.jit('void(f8[:, :], f8[:, :], f8[:, :])')
def cu_square_matrix_mul_shared(A, B, C):
    sA = cuda.shared.array(shape=(tpb, tpb), dtype=float64)
    sB = cuda.shared.array(shape=(tpb, tpb), dtype=float64)

    tx = cuda.threadIdx.x
    ty = cuda.threadIdx.y
    bx = cuda.blockIdx.x
    by = cuda.blockIdx.y
    bw = cuda.blockDim.x
    bh = cuda.blockDim.y

    x = tx + bx * bw
    y = ty + by * bh

    acc = 0.
    for i in range(bpg):
        if x < n and y < n:
            sA[ty, tx] = A[y, tx + i * tpb]
            sB[ty, tx] = B[ty + i * tpb, x]

        cuda.syncthreads()

        if x < n and y < n:
            for j in range(tpb):
                acc += sA[ty, j] * sB[j, tx]

        cuda.syncthreads()

    if x < n and y < n:
        C[y, x] = acc
        
#@cuda.jit('void(f4[:, :], f4[:, :], f4[:, :])')
@cuda.jit('void(f8[:, :], f8[:, :], f8[:, :])')
def cu_square_matrix_mul(A, B, C):
    M = A.shape[0]
    K = A.shape[1]
    N = B.shape[1]
    i, j = cuda.grid(2)
    
    if i < M and j < N:
        acc = 0.
        for k in range(K):
            acc += A[i, k] * B[k, j]
        C[i, j] = acc
        
if __name__ == '__main__':
    M = n
    K = n
    N = n
    blocknum = 16
    blockdim = blocknum, blocknum
    dtype = numpy.float64
    griddim = int(N / blocknum), int(N / blocknum)
    print 'blockdim:', blockdim
    print 'griddim:', griddim   
    
    
    t = time.time()
    numpy.random.seed(0)
    A = numpy.random.random((M, K)).astype(dtype)
    B = numpy.random.random((K, N)).astype(dtype)
    C = numpy.empty((M, N)).astype(dtype)
    cu_square_matrix_mul[griddim, blockdim](A, B, C)
    print "cu_square_matrix_mul        :", time.time() - t
    #print numpy.allclose(C, numpy.dot(A, B))
    
    t = time.time()
    numpy.random.seed(0)
    A = numpy.random.random((M, K)).astype(dtype)
    B = numpy.random.random((K, N)).astype(dtype)
    C = numpy.empty((M, N)).astype(dtype)
    cu_square_matrix_mul_shared[griddim, blockdim](A, B, C)
    print "cu_square_matrix_mul_shared :", time.time() - t
    #print numpy.allclose(C, numpy.dot(A, B))

結果は

blockdim: (16, 16)
griddim: (50, 50)
cu_square_matrix_mul        : 0.359000205994 #Shared Memory なし
cu_square_matrix_mul_shared : 0.077999830246 #Shared Memory あり

約5倍の速度向上となった。
この速度向上なら試す価値はありそうだが、メモリの使い方が職人芸過ぎて浸透しない気がする。。。

行列積をやりたいだけなら、 NVIDIAがcuBLASを提供しているのでそちらが良さそう。

Numbaの有償版のNumbaProにはcuBLASのAPIが提供されているので、お金を払うのも1手です。

CUDA Libraries Host API — Continuum documentation

単精度を試してみる

上のコードは倍精度の計算になっているが、昔はGPUは単精度計算のみだったので速度に寄与するか見てみる。

blockdim: (16, 16)
griddim: (50, 50)
cu_square_matrix_mul        : 0.34299993515   #Shared Memory なし
cu_square_matrix_mul_shared : 0.0940001010895 #Shared Memory あり

誤差レベルと言った感じ。
ただGPUのメモリは1GB程度なことが多いので、
メモリ使用量の面で単精度で良い場合は単精度が良いと思います

最後に

総じてNumbaはPythonの高速化には結構良さげです。

特にCPU利用の場合は、
Cythonほどコードをイジらずに数百倍の高速化が見込めるのは大きいと思います。

一方、GPU利用は、労力の割に益が少ない印象です。
コードもGPU専用に書くので当初の目的からは今ひとつかなと・・・

ただし、最大2兆並列は特定の処理では効果絶大だと思いますので、
知っていて良い知識であることは確かです。

さらにcuBLASのAPIを使うのは良さ気なので$130払って、numpyの高速化を目指すのは良いと思います。
numpyのMKLビルドも付いて来るので悪くないです。

これにてNumba入門は終了です。この後しばらくしたら応用編で、私が勉強するモチベとなった問題を話せればと思います。