Numba pour Python

Passez à la vitesse supérieure en optimisant à la volée votre code Python grâce au compilateur Numba

Numba est un compilateur dont le rôle va être de transformer votre code Python en code machine optimisé. Cette opération est transparente puisque réalisée à la volée pendant l'exécution du code. Numba est cependant principalement dédié aux fonctions mathématiques et de manipulation de tableaux. Déroulons rapidement un exemple pour constater le gain réalisé.

Fonction témoin

Initialisons tout d'abord un tableau numpy à 2 dimensions :

import numpy as np
X = np.random.random((10000,3))

Nous implémentons ensuite une fonction chargée de parcourir ce tableau et de cumuler le carré de chacune des cellules :

def sum_sqrt(X):
  M = X.shape[0]
  N = X.shape[1]
  d=0.0
  for i in range(M):
    for j in range(N):
      d += np.sqrt(X[i,j])
  return d

Ceci etant fait, appelons notre fonction et regardons le temps d'exécution :

time1 = time.perf_counter()
sum_sqrt(X)
print('time: {}'.format(time.perf_counter() - time1))

time: 0.05825133600001209

Fonction Numba

Comme vous le savez Python est un langage interprété, il peut néanmoins être compilé en code machine. L'API Python/C permet au code machine de pouvoir être appelé depuis un programme Python ou l'interpréteur CPython. Nous apportons cette précision car Numba permet via une option noPython de se passer du pont vers Python/C, ce qui amène un gain de temps supplémentaire.
L'utilisation de Numba nécessite bien-sûr l'import de la librairie. Dans notre cas, nous nous contenterons d'importer le package jit.

from numba import jit

Nous réécrivons ensuite la même fonction que précédemment sous un autre nom, en prenant soin d'ajouter un décorateur qui signifiera à Numba que cette fonction est ciblée pour optimisation. Par ailleurs nous spécifions l'option noPython=True.

@jit(nopython=True)
def sum_numba(X):
  M = X.shape[0]
  N = X.shape[1]
  d=0.0
  for i in range(M):
    for j in range(N):
      d += np.sqrt(X[i,j])
  return d

Comme précédemment nous lançons la fonction puis constatons le temps d'exécution.

time1 = time.perf_counter()
sum_numba(X)
print('time: {}'.format(time.perf_counter() - time1))

time: 0.00021922100000892897

Conclusion

Nous passons, grâce au compilateur Numba, de 58.2 ms à 0.22 ms, ce qui vous en conviendrez est une performance intéressante. Numba ne se limite, bien entendu, pas à sa fonction jit, même si il s'agit-la de la plus connue. Nous vous invitons à visiter la documentation officielle pour plus d'information. (doc Numba)
Retrouvez dans la rubrique "Nos datasets" toutes les données dont vous aurez besoin pour tester et pratiquer !