Numba#

Qué es Numba#

Numba, al igual que Cython, es una herramienta para acelerar la ejecución de código escrito en lenguaje Python. Numba puede ser una muy buena alternativa al uso de «Python bindings» con C como ctypes (mecanismo nativo de Python), CFFI o PyBind11; o con Fortran con F2Py (NumPy). No hace falta que aprendas a programar en C o Fortran, con tus rutinas en Python puedes acercarte a tiempos de ejecución con Numba similares a la obtenida con lenguajes compilados.

Numba es una iniciativa muy bien soportada y con estabilidad asegurada a largo plazo. El proyecto lo comenzó en 2012 Travis Oliphant, fundador de Anaconda y creador de Numpy (2005). Actualmente, y desde 2018, Numba está desarrollado por Anaconda y cuenta con el soporte de agencias gubernamentales estado unidenses y empresas tan relevantes como DARPA, Nvidia, Intel o AMD.

Técnicamente, podemos describir a Numba como un compilador Just in time -JIT- de código abierto que traduce código Python y Numpy a código máquina. Esto se lleva a cabo usando LLVM gracias a la librería de Python llvmlite.

Cómo se instala#

Numba puede ser instalado mediante el gestor de paquetes y entornos Conda. Con el entorno en el que se quiere instalar la librería activado, basta con ejecutar:

conda install -c conda-forge numba

O si estás haciendo uso de mamba para poder realizar instalaciones y actualizaciones de tus librerías de manera más rápida:

mamba install -c conda-forge numba

Cómo se usa#

Para usar conda no hace falta que reescribamos ni alteremos sustancialmente nuestro código en Python. Basta con importar la librería y usar adecuadamente el decorador que indicará la compilación de una función, por ejemplo.

Tomemos como ejemplo la siguiente función:

import numpy as np

def sort(serie):
    
    n_numbers = serie.shape[0]
    output = serie.copy()
    
    done = False
    
    while not done:
        done = True
        for ii in range(n_numbers-1):
            if output[ii]>output[ii+1]:
                output[ii], output[ii+1] = output[ii+1], output[ii]
                done = False
    
    return output

Probemos a ordenar una serie de por ejemplo \(10^5\) números enteros positivos aleatoriamente generados en el rango \([0,10^6)\):

rng = np.random.default_rng()

serie = rng.integers(0, 10**5, 10**4)

%%time
sorted_serie = sort(serie)

CPU times: user 27 s, sys: 35 ms, total: 27 s
Wall time: 27 s

print(sorted_serie)

[   22    27    45 ... 99979 99983 99989]

Python no es un interpretador muy rápido y cuando se trata de cómputo pesado, Python puede ser una muy mala opción. Veamos qué pasa cuando le pedimos a Numba que compile la función:

import numba as nb

@nb.jit(nopython=True)
def sort(serie):
    
    n_numbers = serie.shape[0]
    output = serie.copy()
    
    done = False
    
    while not done:
        done = True
        for ii in range(n_numbers-1):
            if output[ii]>output[ii+1]:
                output[ii], output[ii+1] = output[ii+1], output[ii]
                done = False
    
    return output

El decorador jit con la opción nopython=True es invocado para compilar completamente la función sin dejar ninguna linea o fragmento para ser ejecutada en puro Python. Nótese que @nb.jit(nopython=True) es equivalente a @nb.njit() y por lo tanto cualquiera de las dos fórmulas pudo haber sido usada con el mismo fin.

Vamos ahora en cuanto tiempo ordena nuestra función compilada la misma serie de números enteros empleada anteriormente:

%%time
sorted_serie = sort(serie)

CPU times: user 477 ms, sys: 8 ms, total: 485 ms
Wall time: 485 ms

La primera vez que la función se ejecuta, Numba compila de acuerdo al tipo de variables de entrada y salida el código. Esto se puede evitar como veremos más adelante. Comparemos el tiempo tardado con el tiempo necesario en una segunda o tercera vez:

%%time
sorted_serie = sort(serie)

CPU times: user 137 ms, sys: 0 ns, total: 137 ms
Wall time: 137 ms

%%time
sorted_serie = sort(serie)

CPU times: user 137 ms, sys: 4 µs, total: 137 ms
Wall time: 137 ms

Numba ha acelerado la ejecución de nuestra función en un factor mayor a \(150\) veces. Veamos ahora un par de consejos que harán que comiences a usar Numba en tu librería de una manera más eficiente antes de que inviertas un poco de tiempo revisando su documentación.

Es conveniente aportar el tipado de los argumentos de entrada y de los objetos de salida#

En lenguaje de Numba, la firma (signature) es la lista de tipos de de los argumentos de argumentos de entrada junto con la lista de tipos de los argumentos de salida de una función. Aportar explícitamente la firma le va a ayudar a Numba a ser más eficiente compilando el código. Probemos a hacerlo:

import numba as nb

@nb.jit(nb.int64[:](nb.int64[:]), nopython=True)
def sort(serie):
    
    n_numbers = serie.shape[0]
    output = serie.copy()
    
    done = False
    
    while not done:
        done = True
        for ii in range(n_numbers-1):
            if output[ii]>output[ii+1]:
                output[ii], output[ii+1] = output[ii+1], output[ii]
                done = False
    
    return output

%%time
sorted_serie = sort(serie)

CPU times: user 133 ms, sys: 4.01 ms, total: 137 ms
Wall time: 137 ms

Como verás, los tiempos obtenidos la primera vez tras la compilación de la función son similares a la segunda o tercera vez cuando la firma no es proporcionada.

Evita el compilado del código cada vez que tu librería es importada#

Si vas a definir funciones en tu librería que van a ser compiladas con Numba, siguiendo las indicaciones anteriores Numba va a compilar el código cada vez que la librería es importada. Y esto va a suceder incluso si el código de la librería no ha cambiado en nada. Para evitar esta situación puedes hacer uso de la opción del decorador cache=True.

import numba as nb

@nb.jit(nb.int64[:](nb.int64[:]), nopython=True, cache=True)
def sort(serie):
    
    n_numbers = serie.shape[0]
    output = serie.copy()
    
    done = False
    
    while not done:
        done = True
        for ii in range(n_numbers-1):
            if output[ii]>output[ii+1]:
                output[ii], output[ii+1] = output[ii+1], output[ii]
                done = False
    
    return output

De esta manera Numba sólo compilará la función de tu librería la primera vez que ésta es importada. Y mientras el código fuente de la función no sea modificado, no volverá a compilar la función. Verás entonces que el tiempo que cuesta importar tu librería la primera vez es mayor que las veces siguientes. Un precio que se paga gustósamente para poder alcanzar tiempos de ejecución varios órdenes de magnitud más pequeños que si estuvieramos interpretando puro Python.

https://nyu-cds.github.io/python-numba/

Numba

Contenido