web-dev-qa-db-ja.com

numba @jitですべてのCPUコアを使用する方法（numba @jitを並列化）

Pythonで2つのnumpy配列を追加するためにnumbas _@jit_デコレーターを使用しています。 pythonと比較して_@jit_を使用すると、パフォーマンスが非常に高くなります。

ただし、@numba.jit(nopython = True, parallel = True, nogil = True)を渡してもすべてのCPUコアを利用しているわけではありませんです。

Numba _@jit_ですべてのCPUコアを利用する方法はありますか？.

これが私のコードです：

_import time                                                
import numpy as np                                         
import numba                                               

SIZE = 2147483648 * 6                                      

a = np.full(SIZE, 1, dtype = np.int32)                     

b = np.full(SIZE, 1, dtype = np.int32)                     

c = np.ndarray(SIZE, dtype = np.int32)                     

@numba.jit(nopython = True, parallel = True, nogil = True) 
def add(a, b, c):                                          
    for i in range(SIZE):                                  
        c[i] = a[i] + b[i]                                 

start = time.time()                                        
add(a, b, c)                                               
end = time.time()                                          

print(end - start)                                        
_
pythonnumpyjitmulticorenumba
17
user8353052

parallel=Trueを任意のnumba jitted関数に渡すことができますが、それは常にすべてのコアを利用しているという意味ではありません。 numbaはいくつかのヒューリスティックを使用してコードを並列実行させることを理解する必要があります。これらのヒューリスティックは単にコード内で並列化するものを見つけられない場合があります。現在 pull request があるため、「並列」にすることができなかった場合に警告が発行されます。したがって、「並列実行を強制する」ではなく、「可能であれば並列で実行させてください」パラメータに似ています。

ただし、コードを並列化できることが本当にわかっている場合は、常に手動でスレッドまたはプロセスを使用できます。 numba docsからマルチスレッドを使用する例 を変更するだけです。

#!/usr/bin/env python
from __future__ import print_function, division, absolute_import

import math
import threading
from timeit import repeat

import numpy as np
from numba import jit

nthreads = 4
size = 10**7  # CHANGED

# CHANGED
def func_np(a, b):
    """
    Control function using Numpy.
    """
    return a + b

# CHANGED
@jit('void(double[:], double[:], double[:])', nopython=True, nogil=True)
def inner_func_nb(result, a, b):
    """
    Function under test.
    """
    for i in range(len(result)):
        result[i] = a[i] + b[i]

def timefunc(correct, s, func, *args, **kwargs):
    """
    Benchmark *func* and print out its runtime.
    """
    print(s.ljust(20), end=" ")
    # Make sure the function is compiled before we start the benchmark
    res = func(*args, **kwargs)
    if correct is not None:
        assert np.allclose(res, correct), (res, correct)
    # time it
    print('{:>5.0f} ms'.format(min(repeat(lambda: func(*args, **kwargs),
                                          number=5, repeat=2)) * 1000))
    return res

def make_singlethread(inner_func):
    """
    Run the given function inside a single thread.
    """
    def func(*args):
        length = len(args[0])
        result = np.empty(length, dtype=np.float64)
        inner_func(result, *args)
        return result
    return func

def make_multithread(inner_func, numthreads):
    """
    Run the given function inside *numthreads* threads, splitting its
    arguments into equal-sized chunks.
    """
    def func_mt(*args):
        length = len(args[0])
        result = np.empty(length, dtype=np.float64)
        args = (result,) + args
        chunklen = (length + numthreads - 1) // numthreads
        # Create argument tuples for each input chunk
        chunks = [[arg[i * chunklen:(i + 1) * chunklen] for arg in args]
                  for i in range(numthreads)]
        # Spawn one thread per chunk
        threads = [threading.Thread(target=inner_func, args=chunk)
                   for chunk in chunks]
        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join()
        return result
    return func_mt


func_nb = make_singlethread(inner_func_nb)
func_nb_mt = make_multithread(inner_func_nb, nthreads)

a = np.random.Rand(size)
b = np.random.Rand(size)

correct = timefunc(None, "numpy (1 thread)", func_np, a, b)
timefunc(correct, "numba (1 thread)", func_nb, a, b)
timefunc(correct, "numba (%d threads)" % nthreads, func_nb_mt, a, b)

変更した部分を強調表示し、その他はすべて例からそのままコピーしました。これは私のマシンのすべてのコア（4コアマシン、したがって4スレッド）を利用しますが、大幅なスピードアップは見られません。

numpy (1 thread)       539 ms
numba (1 thread)       536 ms
numba (4 threads)      442 ms

この場合のマルチスレッディングによる（大幅な）高速化の欠如は、追加が帯域幅制限のある操作であることです。つまり、実際の加算を行うよりも、配列から要素をロードして結果配列に結果を配置する方がはるかに時間がかかります。

これらの場合、並列実行が原因で速度が低下することさえあります！

関数がより複雑で、実際の操作が配列要素の読み込みと保存に比べてかなり時間がかかる場合にのみ、並列実行で大きな改善が見られます。 numbaのドキュメントの例は、次のようなものです。

def func_np(a, b):
    """
    Control function using Numpy.
    """
    return np.exp(2.1 * a + 3.2 * b)

@jit('void(double[:], double[:], double[:])', nopython=True, nogil=True)
def inner_func_nb(result, a, b):
    """
    Function under test.
    """
    for i in range(len(result)):
        result[i] = math.exp(2.1 * a[i] + 3.2 * b[i])

2つの乗算、1つの加算、およびmath.expの1回の呼び出しは、結果のロードと格納よりもはるかに遅いため、これは実際にはスレッドの数に応じて（ほぼ）スケーリングされます。

func_nb = make_singlethread(inner_func_nb)
func_nb_mt2 = make_multithread(inner_func_nb, 2)
func_nb_mt3 = make_multithread(inner_func_nb, 3)
func_nb_mt4 = make_multithread(inner_func_nb, 4)

a = np.random.Rand(size)
b = np.random.Rand(size)

correct = timefunc(None, "numpy (1 thread)", func_np, a, b)
timefunc(correct, "numba (1 thread)", func_nb, a, b)
timefunc(correct, "numba (2 threads)", func_nb_mt2, a, b)
timefunc(correct, "numba (3 threads)", func_nb_mt3, a, b)
timefunc(correct, "numba (4 threads)", func_nb_mt4, a, b)

結果：

numpy (1 thread)      3422 ms
numba (1 thread)      2959 ms
numba (2 threads)     1555 ms
numba (3 threads)     1080 ms
numba (4 threads)      797 ms
18
MSeifert

完全を期すために、2018年（numba v 0.39）には、

from numba import prange

そして、元の関数定義でrangeをprangeに置き換えます。それだけです。

これにより、すぐにCPU使用率が100％になり、私の場合はランタイムの2.9秒から1.7秒にスピードアップします（16コア32スレッドのマシンでSIZE = 2147483648 * 1の場合）。

より複雑なカーネルの場合、fastmath=Trueを渡すことで、さらに高速化できます。

7
Anatoly Alekseev