Numbaを使ってpyhonコードを高速化する方法

こんにちは、Link-Uの町屋敷です。

今回は、機械学習を行う際にほぼ必ず行わなければならない前処理を、
GPUを使ってやったら早く終わったので、メモがてらに書いていきます。
今回はpythonでの話です。

pythonでcudaを使ったGPU演算をする方法として、
NVIDIAの公式でも紹介されているnumbaを使いたいと思います。

@Vectorizeを使った方法

pipでも入るそうですが、anacondaでもインストールできるのでcondaを使用します。

Anaconda,Minicondaのインストール方法はバックナンバーで、

conda install accelerate

このコマンドを使うとnamba以外にも必要なパッケージが入るので楽です。

早速試してみましょう。

先程の公式ページの動画内に書かれていたコードを使っててみます。

import numpy as np

from timeit import default_timer as timer
from numba import vectorize
from numba import cuda, float32

@vectorize(['float32(float32, float32)'], target='cpu')
def VectorAdd(a,b):
    return a + b

@vectorize(['float32(float32, float32)'], target='cuda')
def GpuAdd(a,b):
    return a + b

def NormalAdd(a,b,c):
    for i, _ in enumerate(a):
        c[i] = b[i] + a[i]    

def main():
    
    for N in [1e2, 1e3, 1e4, 1e5, 1e6, 1e7, 1e8]:
        vec_a = np.ones(int(N), dtype=np.float32)
        vec_b = np.ones(int(N), dtype=np.float32)
        vec_c = np.zeros(int(N), dtype=np.float32)
        
        start = timer()
        NormalAdd(vec_a, vec_b, vec_c)
        normal_time = timer() - start
        
        print('When [{0}] vector, NormalAdd took {1} seconds'.format(N, normal_time))
        
        vec_c = np.zeros(int(N), dtype=np.float32)
        
        start = timer()
        vec_c = VectorAdd(vec_a, vec_b)
        cpu_time = timer() - start
        #print(C)
        
        print('When [{0}] vector, CpuAdd took {1} seconds'.format(N, cpu_time))
        
        start = timer()
        vec_c = GpuAdd(vec_a, vec_b)
        gpu_time = timer() - start
        
        print('When [{0}] vector, GpuAdd took {1} seconds'.format(N, gpu_time))
        print('Normal speed / CPU speed = {0}'.format(normal_time/cpu_time))       
        print('Normal speed / GPU speed = {0}'.format(normal_time/gpu_time))    
        print('CPU speed / GPU speed = {0}'.format(cpu_time/gpu_time))

numbaでは高速化を行いたい関数にデコレータを付けます。

この例では@vectorize([‘float32(float32, float32)’], target=’cuda’)の部分が該当します。

@vectorizeの第一引数には、C言語の関数の宣言ように返り値と引数の型を記述します。

第二引数はをcudaにするとGPUをcpuにするとcpuを使って最適化します。

GPUサーバーで動かした結果はこの通り。

何もしていないときと比べて遥かに早くなりました。

ただ、GPUを使用し始めるのに3.5秒ほど準備期間が必要な模様で、最初の処理に時間がかかっています。

また、CPUのほうが早いという結果に。

桁数がもう少し増えれば逆転しそうだがあまり降らしすぎるとメモリーエラーになります。

公式の動画でtarget=cpuをやっていないのはそういうことがややこしいからなのか……

このままではCPUのほうが良いんじゃないか説が出てきてしまうので、少し処理を追加した関数でやってみましょう。

@vectorize(['float32(float32)'], target='cuda')
def RadToDegGpu(a):
    return (360 + 180 * a / np.pi) % 360

@vectorize(['float32(float32)'], target='cpu')
def RadToDegCpu(a):
    return (360 + 180 * a / np.pi) % 360

def main():
    
    for N in [1e2, 1e3, 1e4, 1e5, 1e6, 1e7, 1e8]:
        A = np.ones(int(N), dtype=np.float32) * np.pi
        B = np.zeros(int(N), dtype=np.float32)
        
        start = timer()
        B = RadToDegCpu(A)
        cpu_time = timer() - start
        
        print('When [{0}] vector, RadToDegCpu took {1} seconds'.format(N, cpu_time))
        
        start = timer()
        B = RadToDegGpu(A)
        gpu_time = timer() - start
         
        print('When [{0}] vector, RadToDegGpu took {1} seconds'.format(N, gpu_time))
        print('x{0} speed'.format(cpu_time/gpu_time))

ラジアンを°に変える簡単な関数です。

nambaはnumpyに最適化されているため関数を使う時は極力numpyを使いましょう(この例だと使ってるのはただの定数だけど)。

引数の数が変わっているので、@vectorizeの第一引数も変わっています。

x~speedの~が1以上だとgpuのほうが早いです。

桁数が小さい時はまだCPUのほうが早いですが、大きくなってくるとGPUを使ったほうが約3倍早くなっています。

ただし最初の準備期間3.5秒は一律でかかるので、この処理単体だけをやるとGPUのほうが遅くなります。

1回の実行では使う関数が変わっても1回しか準備期間はないようなので、封数の処理を一気にやると良いでしょう。

@cuda.jitを使った方法

次に行列を使った計算を高速化する方法をここを参考にして書いていきます。

例としてデータの前処理によく使う各行に対して平滑化を行うプログラムを作ります。

まず、普通の実装

def NormalSmooth(inputMat, outputMat):
    for i, row in enumerate(inputMat):
        for j, v in enumerate(row):
            if j == 0:
                outputMat[i, j] = (inputMat[i, j] + inputMat[i, j] + inputMat[i, j + 1])/3
            elif j == outputMat.shape[1] - 1:
                outputMat[i, j] = (inputMat[i, j - 1] + inputMat[i, j] + inputMat[i, j])/3
            else:
                outputMat[i, j] = (inputMat[i, j - 1] + inputMat[i, j] + inputMat[i, j + 1])/3

そして、@cuda.jitを使った実装

@cuda.jit
def GpuSmooth(inputMat, outputMat):
    row, col = cuda.grid(2)
    if row < outputMat.shape[0] and col < outputMat.shape[1]:
        if col == 0:
            outputMat[row, col] = (inputMat[row, col] + inputMat[row, col] + inputMat[row, col + 1])/3
        elif col == outputMat.shape[1] - 1:
            outputMat[row, col] = (inputMat[row, col - 1] + inputMat[row, col] + inputMat[row, col])/3
        else:
            outputMat[row, col] = (inputMat[row, col - 1] + inputMat[row, col] + inputMat[row, col + 1])/3

cuda.grid(2)を使うことで簡単に2重のfor文の処理を書き換えることが出来ます、

cuda.grid(2)が返す値はx,yともに正の値ですが、上限が行列の形と違うのでif文を使って行列の外になった場合処理をしないことが必要になります。

def main(): 
     
    for R in [10,100,1000,10000]:
        for C in [10,100,1000,10000]:
    
            A = np.ones((R, C))
            B = np.ones((R, C))
            
            start = timer()
            NormalSmooth(B, A)
            normal_time = timer() - start
            
            print('When [{0}, {1}] matrix, NormalSmooth took {2} seconds'.format(R, C, normal_time))
            
            start = timer()
            GpuSmooth(B, A)
            gpu_time = timer() - start
            
            print('When [{0}, {1}] matrix, GpuSmooth took {2} seconds'.format(R, C, gpu_time))
            print('x{0} speed'.format(normal_time/gpu_time))

こちらも行列が小さい時はCPUのほうが優位ですが大きくなってくると最大で129倍と圧倒的にGPUのほうが早くなります。

ちなみに手元のノートパソコンの結果が[10000, 10000]のとき121.75秒かかっていたので100倍以上高速化されたことになります。

まとめ

numbaを使うことでpytonスクリプトを高速化する手法を書いた

大きいデータを扱う時はかなり高速化されるので積極的に使っていきたい。