La plus rapide des log-sum-exp de Python dans un 'reduceat' (réduction)

Question

Dans le cadre d'une programmation statistique, je dois ajouter des valeurs log-transformées avec la valeur Fonction LogSumExp . Cette méthode est nettement moins efficace que l'addition de valeurs non calculées.

En outre, j'ai besoin d'additionner des valeurs à l'aide de la fonction numpy.ufunc.reduecat fonctionnalité.

J'ai envisagé plusieurs options, dont le code figure ci-dessous :

1. (pour une comparaison dans un espace non logarithmique) utiliser numpy.add.reduceat
2. L'ufunc de Numpy pour additionner des valeurs enregistrées : np.logaddexp.reduceat

3. Fonction de réduction manuscrite avec les fonctions logsumexp suivantes :

   • l'implémentation de logsumexp par scipy
   • fonction logsumexp en Python (avec numba )
   • Fonction logsumexp en continu en Python (avec numba )

   def logsumexp_reduceat(arr, indices, logsum_exp_func): res = list() i_start = indices[0] for cur_index, i in enumerate(indices[1:]): res.append(logsum_exp_func(arr[i_start:i])) i_start = i

   res.append(logsum_exp_func(arr[i:]))
   return res

   @numba.jit(nopython=True) def logsumexp(X): r = 0.0 for x in X: r += np.exp(x)
   return np.log(r)

   @numba.jit(nopython=True) def logsumexp_stream(X): alpha = -np.Inf r = 0.0 for x in X: if x != -np.Inf: if x <= alpha: r += np.exp(x - alpha) else: r *= np.exp(alpha - x) r += 1.0 alpha = x return np.log(r) + alpha

   arr = np.random.uniform(0,0.1, 10000) log_arr = np.log(arr) indices = sorted(np.random.randint(0, 10000, 100))

   approach 1

   %timeit np.add.reduceat(arr, indices) 12.7 µs ± 503 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

   approach 2

   %timeit np.logaddexp.reduceat(log_arr, indices) 462 µs ± 17.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

   approach 3, scipy function

   %timeit logsum_exp_reduceat(arr, indices, scipy.special.logsumexp) 3.69 ms ± 273 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

   approach 3 handwritten logsumexp

   %timeit logsumexp_reduceat(log_arr, indices, logsumexp) 139 µs ± 7.1 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

   approach 3 streaming logsumexp

   %timeit logsumexp_reduceat(log_arr, indices, logsumexp_stream) 164 µs ± 10.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Les résultats de timeit montrent que les fonctions logsumexp écrites à la main avec numba sont les options les plus rapides, mais sont toujours 10x plus lentes que numpy.add.reduceat.

Quelques questions :

1. Existe-t-il d'autres approches (ou des modifications des options que j'ai présentées) qui soient plus rapides ? Par exemple, existe-t-il un moyen d'utiliser une table de recherche pour calculer la fonction logsumexp ?
2. Pourquoi la fonction "streaming logsumexp" de Sebastian Nowozin n'est-elle pas plus rapide que l'approche naïve ?

Answer 1

Il y a une certaine marge d'amélioration

Mais ne vous attendez jamais à ce que les logsumexp soient aussi rapides qu'une sommation standard, car exp est une opération assez coûteuse.

Exemple

import numpy as np

#from version 0.43 until 0.47 this has to be set before importing numba
#Bug: https://github.com/numba/numba/issues/4689
from llvmlite import binding
binding.set_option('SVML', '-vector-library=SVML')
import numba as nb

@nb.njit(fastmath=True,parallel=False)
def logsum_exp_reduceat(arr, indices):
    res = np.empty(indices.shape[0],dtype=arr.dtype)

    for i in nb.prange(indices.shape[0]-1):
        r = 0.
        for j in range(indices[i],indices[i+1]):
            r += np.exp(arr[j])  
        res[i]=np.log(r)

    r = 0.
    for j in range(indices[-1],arr.shape[0]):
        r += np.exp(arr[j])  
    res[-1]=np.log(r)
    return res

Horaires

#small example where parallelization doesn't make sense
arr = np.random.uniform(0,0.1, 10_000)
log_arr = np.log(arr)
#use arrays if possible
indices = np.sort(np.random.randint(0, 10_000, 100))

%timeit logsum_exp_reduceat(arr, indices)
#without parallelzation 22 µs ± 173 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
#with parallelization   84.7 µs ± 32.2 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit np.add.reduceat(arr, indices)
#4.46 µs ± 61.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

#large example where parallelization makes sense
arr = np.random.uniform(0,0.1, 1000_000)
log_arr = np.log(arr)
indices = np.sort(np.random.randint(0, 1000_000, 100))

%timeit logsum_exp_reduceat(arr, indices)
#without parallelzation 1.57 ms ± 14.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
#with parallelization   409 µs ± 14.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.add.reduceat(arr, indices)
#340 µs ± 11.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)