Parfois, vous devez écrire du code numpy non idiomatique si vous voulez vraiment accélérer votre calcul, ce que vous ne pouvez pas faire avec numpy natif.
numba
compile votre code python en bas niveau C. Étant donné que beaucoup de numpy lui-même est généralement aussi rapide que C, cela finit surtout par être utile si votre problème ne se prête pas à la vectorisation native avec numpy. Ceci est un exemple (où j'ai supposé que les indices sont contigus et triés, ce qui se reflète également dans les données d'exemple):
import numpy as np
import numba
# use the inflated example of roganjosh https://stackoverflow.com/a/58788534
data = [1.00, 1.05, 1.30, 1.20, 1.06, 1.54, 1.33, 1.87, 1.67]
index = [0, 0, 1, 1, 1, 1, 2, 3, 3]
data = np.array(data * 500) # using arrays is important for numba!
index = np.sort(np.random.randint(0, 30, 4500))
# jit-decorate; original is available as .py_func attribute
@numba.njit('f8[:](f8[:], i8[:])') # explicit signature implies ahead-of-time compile
def diffmedian_jit(data, index):
res = np.empty_like(data)
i_start = 0
for i in range(1, index.size):
if index[i] == index[i_start]:
continue
# here: i is the first _next_ index
inds = slice(i_start, i) # i_start:i slice
res[inds] = data[inds] - np.median(data[inds])
i_start = i
# also fix last label
res[i_start:] = data[i_start:] - np.median(data[i_start:])
return res
Et voici quelques synchronisations utilisant la %timeit
magie d'IPython :
>>> %timeit diffmedian_jit.py_func(data, index) # non-jitted function
... %timeit diffmedian_jit(data, index) # jitted function
...
4.27 ms ± 109 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
65.2 µs ± 1.01 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
En utilisant les données d'exemple mises à jour dans la question, ces nombres (c'est-à-dire le temps d'exécution de la fonction python par rapport au temps d'exécution de la fonction accélérée JIT) sont
>>> %timeit diffmedian_jit.py_func(data, groups)
... %timeit diffmedian_jit(data, groups)
2.45 s ± 34.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
93.6 ms ± 518 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
Cela équivaut à une accélération de 65x dans le petit cas et une accélération de 26x dans le plus grand cas (par rapport au code en boucle lente, bien sûr) en utilisant le code accéléré. Un autre avantage est que (contrairement à la vectorisation typique avec numpy natif), nous n'avions pas besoin de mémoire supplémentaire pour atteindre cette vitesse, il s'agit de code de bas niveau optimisé et compilé qui finit par être exécuté.
La fonction ci-dessus suppose que les tableaux numpy int sont int64
par défaut, ce qui n'est pas réellement le cas sous Windows. Une alternative consiste donc à supprimer la signature de l'appel à numba.njit
, déclenchant une compilation juste à temps appropriée. Mais cela signifie que la fonction sera compilée lors de la première exécution, ce qui peut interférer avec les résultats de synchronisation (nous pouvons soit exécuter la fonction une fois manuellement, en utilisant des types de données représentatifs, soit simplement accepter que la première exécution de synchronisation soit beaucoup plus lente, ce qui devrait Etre ignoré). C'est exactement ce que j'ai essayé d'empêcher en spécifiant une signature, ce qui déclenche une compilation anticipée.
Quoi qu'il en soit, dans le cas JIT correctement, le décorateur dont nous avons besoin est juste
@numba.njit
def diffmedian_jit(...):
Notez que les timings ci-dessus que j'ai montrés pour la fonction compilée jit ne s'appliquent qu'une fois la fonction compilée. Cela se produit soit lors de la définition (avec une compilation désirée, lorsqu'une signature explicite est passée à numba.njit
), soit pendant le premier appel de fonction (avec une compilation différée, lorsqu'aucune signature n'est transmise à numba.njit
). Si la fonction ne doit être exécutée qu'une seule fois, le temps de compilation doit également être pris en compte pour la vitesse de cette méthode. Cela ne vaut généralement la peine de compiler des fonctions que si le temps total de compilation + exécution est inférieur au temps d'exécution non compilé (ce qui est en fait vrai dans le cas ci-dessus, où la fonction python native est très lente). Cela se produit principalement lorsque vous appelez souvent votre fonction compilée.
Comme le note max9111 dans un commentaire, une des caractéristiques importantes de numba
est le cache
mot - clé to jit
. Passer cache=True
à numba.jit
stockera la fonction compilée sur le disque, de sorte que lors de la prochaine exécution du module python donné, la fonction sera chargée à partir de là plutôt que recompilée, ce qui peut encore vous épargner l'exécution à long terme.
scipy.ndimage.median
suggestion dans la réponse liée? Il ne me semble pas qu'il ait besoin d'un nombre égal d'éléments par étiquette. Ou ai-je raté quelque chose?