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スクリプトを高速化する手法を書いた

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

AIを端末側で処理するチップ達

以前NVIDIAの新しいGPUを紹介しましたが、NVIDIAのGPUはPCで主に利用されると思います。

しかし、恐らくAI(正確には機械学習の推論)を行う回数で言うと圧倒的にスマートフォンの方が多いと思います。

今回はスマートフォン等に搭載されるチップがAIをどう処理するかを紹介していきましょう。今回はどのチップもPVがあったので全編PV付きでお送りします。

Snapdragon 845

スマホ用チップの雄、QUALCOMMの最上位チップです。Androidのハイエンド機は一部の例外を除いてほぼこのチップを使っていると言っても過言ではないかと。

ただし、最近のチップの中では比較的AI系の性能は控えめで、どちらかというと従来からのCPUやGPUといった部分の強い正統派のチップです。

AIに対する対応はCPU+GPU+DSPで行い専用回路は持ちません。

この対応だと、AI用の回路がない分、チップに余裕ができるので、CPUやGPUといった汎用的な回路を強化できる利点があります。

一方特にCPUなどは汎用的な処理ができる分、大量のAI向け演算などは処理量の点でも、電力効率の点でも苦手と言えます。

このあたりAIに関連する処理が、全スマホ利用時間のうちどの程度を占めるか、メーカーによる考え方の違いとなって表れてきます。

QUALCOMMとしては音声認識など局所的にAIを使うことはあっても、全利用時間に占める割合はさほど高くないと踏んでいるのではないでしょうか。

とはいえ、AI性能は全世代のSnapdragon 835比で3倍をうたっていますので、CPUやGPUの性能が30%程度の向上であることを加味すると力を入れてきているなというのはあります。

ところでSnapdragon 845にはWiFiの接続が16倍高速になるという機能向上もあり、性能面で言うとそちらが一番の底上げであったり。

そちらも個人的には気になります。

Kirin 980

最近登場した中国Huawei製のチップ。恐らくHuawei製の端末以外採用されていない・・・はず。

QUALCOMMとは真逆で積極的にAI系の機能を盛り込んでくるメーカーで、昨年のKirin 970からAI系の回路を盛り込んできています。

Kirin 980では回路規模を倍増させており、AI系の性能はKirin 970の2倍以上になっていると思われます。

とはいえ、CPU性能も75%向上したらしく、AI系の性能を特別に向上させた、というよりは全般的に2倍程度になるようにチップを設計したようです。

中国メーカーの技術力の伸びを示すような、かなりアグレッシブな性能のあげ方です。

Apple A12 Bionic

日本で一番普及しているiPhoneの最新版に搭載されているチップ。

Appleは他の会社と違ってちゃんと素の演算回数を公表してくれており毎秒5超回の演算ができます。

全世代からAI系の機能をチップに入れてきていますが、前世代と比較して9倍の性能になっています。

こうして他のチップと比較してみるとAppleが全力でAI系を増強してきていることがわかります。

実はAppleは自社生産かつ大量生産なので、他の会社よりもチップにお金をかけることができます。

なので、チップに他社より多くの機能を入れることができ、AI系の機能を強化する余裕があるのかもしれません。

Movidius NCS

IntelがAIチップベンチャーを買収して作ったUSBデバイス。1Wで100GFlopsです。

専用設計で100GFlopsというのは若干微妙な気もしないではないですが、USBインターフェイスとかも入っていることも考えるとこんなものでしょうか。

このデバイス、USBで刺すと使えるので後付けできるのが最大の魅力です。しかも複数台刺して性能向上なんてこともできるようです。

Raspberry Piに刺して使えるなどD.I.Yでデバイスが作れそうな予感が大変します。

というかIntelもCPUにそういう系の機能つければいいのにと思わなくもないです。

まとめ

各メーカーで取り組み方に差はみられるものの、基本的にどのチップメーカーもAI強化に舵を切っていることがわかります。

現在はカメラ、音声認識などでの利用が多いようですが、AIを使えば便利になるシーンは多いと思います。