Acelerar el código Python con Numba

Posiblemente Python sea uno de nuestros leguajes de programación favoritos debido a su sencillez y potencia. Pero hay que reconocer que es lento al compararlos con otros como pueden ser C, Julia o JavaScript. Por eso existen múltiples soluciones para aumentar el rendimiento del código Python. Algunas que hemos visto en entradas anteriores son el uso de cauces con o la complicación de código con Cython. Otra opción que podemos usar para acelerar el código Python es Numba.

¿Qué es Numba?

Numba es un compilador JIT que permite traducir un subconjunto de funciones Python y NumPy en código máquina usando LLVM. La misma arquitectura que usa por ejemplo Julia. Ofreciendo además la posibilidad de paralizar el código tanto en CPU como GPU casi sin necesidad de cambiar el código.

Numba se instala como cualquier otro paquete de Python, por lo que si queremos utilizarlo solamente tenemos que instalar el paquete con PIP. Para lo necesitamos escribir en la terminal el comando:

pip install numba

Uso de Numba

El uso de Numba es relativamente sencillo. Solamente se tiene que importar njit de numba e incluir este embellecedor antes de la definición de la función. Siendo esto todo lo que hay que hacer para acelerar nuestro código.

Una de las funciones que suelo emplear para medir el rendimiento de código es la sucesión de Fibonacci. Una función sencilla de implementar, pero en la que el tiempo de ejecución crece rápidamente. Así para comprobar cuánto puede mejorar el rendimiento usar Numba se puede escribir el siguiente código.

import timeit
from numba import njit

def fibonacci(n):
    if n < 2:
        return n
    else:
        return fibonacci(n-1) + fibonacci(n-2)
    
@njit(fastmath=True, cache=True)
def fibonacci_numba(n):
    if n < 2:
        return n
    else:
        return fibonacci_numba(n-1) + fibonacci_numba(n-2)
    
py = timeit.timeit("fibonacci(35)", "from __main__ import fibonacci", number=30)
nb = timeit.timeit("fibonacci_numba(35)", "from __main__ import fibonacci_numba", number=30)

En el que solamente se han creado dos funciones idénticas, una de ellas modificada con @njit(fastmath=True, cache=True) para que compile con Numba. Posteriormente se ha utilizado timeit para medir el rendimiento de ambas funciones

Rendimiento de Numba en comparación con Python

Al ejecutar el código anterior se puede ver que el tiempo necesario para ejecutar la función fibonacci es cerca de 35 veces mayor que el tiempo necesario para obtener el resultado con fibonacci_numba. Valor que puede variar en función de la configuración utilizada para probar el código. Un aumento considerable teniendo en cuenta que no ha sido necesario modificar el código de función, simplemente se ha compilado con Numba.

Conclusiones

En esta entrada hemos visto cómo se puede acelerar el código Python con Numba en un factor 35 sin la necesidad de modificar el código. Lo que es una ventaja respecto a otras soluciones como Cython. Quizás en una futura entrada sería una buena opción comparar el rendimiento de Cython, lru_cache y Numba con diferentes tipos de funciones.

Imagen de Michael Kauer en Pixabay