web-dev-qa-db-fra.com

Comment faire pour que numba @jit utilise tous les cœurs du processeur (paralléliser numba @jit)

J'utilise numbas @jit decorator pour ajouter deux tableaux numpy en python. Les performances sont si élevées si j'utilise @jit comparé à python. 

Cependant, il s'agit de n'utilisant pas tous les cœurs du processeur même si je passe en @numba.jit(nopython = True, parallel = True, nogil = True). 

Est-il possible d'utiliser tous les cœurs de processeur avec numba @jit.

Voici mon code:

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
6
user8353052

Vous pouvez passer parallel=True à n’importe quelle fonction numérisée, mais cela ne signifie pas qu’elle utilise toujours tous les cœurs. Vous devez comprendre que numba utilise des heuristiques pour exécuter le code en parallèle. Parfois, ces heuristiques ne trouvent simplement rien à paralléliser dans le code. Il y a actuellement une requête pull afin qu'elle émette un avertissement s'il n'était pas possible de la rendre "parallèle". Donc, cela ressemble plus à un paramètre "s'il vous plaît faites-le exécuter en parallèle si possible" et non pas "appliquer parallèlement".

Cependant, vous pouvez toujours utiliser les threads ou les processus manuellement si vous savez réellement que vous pouvez paralléliser votre code. Il suffit d’adapter le exemple d’utilisation du multi-threading à partir de 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)

J'ai mis en évidence les parties que j'ai modifiées, tout le reste a été copié textuellement de l'exemple. Ceci utilise tous les cœurs de ma machine (4 cœurs donc 4 threads) mais ne montre pas une accélération significative:

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

L'absence (beaucoup) d'accélération avec le multithreading dans ce cas est que l'ajout est une opération limitée en bande passante. Cela signifie qu'il faut beaucoup plus de temps pour charger les éléments du tableau et placer le résultat dans le tableau de résultats que pour effectuer l'ajout proprement dit. 

Dans ces cas, vous pourriez même voir des ralentissements dus à une exécution parallèle!

Si les fonctions sont plus complexes et que l'opération prend un temps considérable par rapport au chargement et au stockage d'éléments de tableau, vous constaterez une amélioration importante avec l'exécution en parallèle. L'exemple dans la documentation de numba en est un comme celui-ci:

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])

Cela correspond en fait (presque) au nombre de threads, car deux multiplications, un ajout et un appel à math.exp est beaucoup plus lent que le chargement et le stockage des résultats:

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)

Résultat:

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
11
MSeifert

Par souci d'exhaustivité, en 2018 (numba v 0.39), vous pouvez simplement faire

from numba import prange

et remplacez range par prange dans la définition de votre fonction d'origine, c'est tout.

Cela rend immédiatement l’utilisation du processeur à 100% et, dans mon cas, accélère les choses de 2 à 1,7 secondes d’exécution (pour SIZE = 2147483648 * 1, sur une machine à 16 cœurs 32 threads).

Les noyaux plus complexes peuvent souvent être encore plus rapides en passant par fastmath=True.

3
Anatoly Alekseev