es.davy.ai

Preguntas y respuestas de programación confiables

¿Tienes una pregunta?

Si tienes alguna pregunta, puedes hacerla a continuación o ingresar lo que estás buscando.

Numba no está mejorando el rendimiento.

27 Jul, 2023 Programación

Estoy probando el rendimiento de numba en una función que toma una matriz numpy y comparo:

import numpy as np
from numba import jit, vectorize, float64
import time
from numba.core.errors import NumbaWarning
import warnings

warnings.simplefilter('ignore', category=NumbaWarning)

@jit(nopython=True, boundscheck=False) # Establezco el modo "nopython" para obtener el mejor rendimiento, equivalente a @njit
def go_fast(a):     # La función se compila en código de máquina cuando se llama por primera vez
    trace = 0.0
    for i in range(a.shape[0]):   # Numba se beneficia de los bucles
        trace += np.tanh(a[i, i]) # Numba se beneficia de las funciones de NumPy
    return a + trace              # Numba se beneficia de la transmisión de NumPy

class Main(object):
    def __init__(self) -> None:
        super().__init__()
        self.mat     = np.arange(100000000, dtype=np.float64).reshape(10000, 10000)

    def my_run(self):
        st = time.time()
        trace = 0.0
        for i in range(self.mat.shape[0]):   
            trace += np.tanh(self.mat[i, i]) 
        res = self.mat + trace
        print('Duración de Python: ', time.time() - st)
        return res                           

    def jit_run(self):
        st = time.time()
        res = go_fast(self.mat)
        print('Duración de Jit: ', time.time() - st)
        return res

obj = Main()
x1 = obj.my_run()
x2 = obj.jit_run()

La salida es:

Duración de Python:  0.2164750099182129
Duración de Jit:  0.5367801189422607

¿Cómo puedo obtener una versión mejorada de este ejemplo?

Tags:  , , ,

Answer

  1. Avatar for davy.ai
    davy.ai

    Para mejorar el rendimiento de este ejemplo, puedes intentar usar el decorador vectorize del módulo numba en lugar del decorador jit. El decorador vectorize te permite crear una ufunc (función universal) que puede operar en matrices de numpy de una manera más rápida.

    Aquí tienes un ejemplo de código que muestra cómo usar este enfoque:

    import numpy as np
from numba import vectorize, float64
import time

@vectorize([float64(float64)], target='parallel')
def go_fast(a):
    return np.tanh(a)

class Main(object):
    def __init__(self):
        super().__init__()
        self.mat = np.arange(100000000, dtype=np.float64).reshape(10000, 10000)

    def my_run(self):
        st = time.time()
        trace = 0.0
        for i in range(self.mat.shape[0]):
            trace += np.tanh(self.mat[i, i])
        res = self.mat + trace
        print('Duración de Python: ', time.time() - st)
        return res

    def jit_run(self):
        st = time.time()
        res = go_fast(self.mat)
        trace = np.sum(go_fast(np.diag(self.mat)))
        res += trace
        print('Duración de Jit: ', time.time() - st)
        return res

obj = Main()
x1 = obj.my_run()
x2 = obj.jit_run()

    En el código anterior, usamos el decorador vectorize para definir la función go_fast como una ufunc que opera en un solo elemento de la matriz de entrada de numpy, devolviendo el valor tanh de ese elemento. Especificamos el tipo de datos de entrada y salida como float64, y establecemos el atributo target en 'parallel' para habilitar la paralelización.

    En el método jit_run, llamamos a la función go_fast en toda la matriz de entrada, y luego calculamos la traza de la matriz usando go_fast en los elementos diagonales. Finalmente, añadimos el valor de la traza a la matriz y devolvemos el resultado.

    Cuando comparamos los tiempos de ejecución de los métodos my_run y jit_run usando este enfoque, deberíamos ver una mejora significativa en el rendimiento.

Comments are closed.