tkm2261's blog

研究員(OR屋) → データ分析官 → MLエンジニア → ニート → 米国CS PhDが諸々書いてます

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で構成しているのは、もったいない気がしていますが、
LinuxにGPU載せても用途が限られすぎるので妥協です。

環境構築

ここまで、特に環境構築について触れて来ませんでしたが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

実際に動かしてみた

その３まで同様に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のメモリは１GB程度なことが多いので、
メモリ使用量の面で単精度で良い場合は単精度が良いと思います

最後に

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

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

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

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

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

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