Si vous faites de l’analyse de gros jeux de données en Python, vous vous êtes vite rendu compte que c’est là que le langage montre sa lenteur.

Les boucles pour manipuler des arrays d’entiers, qui prennent quelque nanosecondes en C, prennent des microsecondes en Python. C’est pour ça que les scientifiques utilisent SciPy et les analystes Pandas, qui sont tous deux basés sur la lib C Numpy bien enrobée dans du joli Python.

Parfois néanmoins, garder le code en pur Python est un objectif, comme dans le cas du projet mss, un screenshotter qui ne souhaites avoir une extension compilée attachée à la patte.

Or la manipulation d’images, ça bouffe des grosses matrices de pixels, et l’auteur nous a posé une jolie colle sur IndexError : ayant un array de pixels en notation BGR (Blue, Green Red), comment le transposer en RGB de manière performante ?

Sa proposition était :

pixels = range(300) # je réduis le nombre de pixel pour le benchmark
buffer_len = len(pixels)
 
def version1():
    for idx in range(0, buffer_len - 2, 3):
        pixels[idx + 2], pixels[idx] = pixels[idx], pixels[idx + 2]

J’ai tenté :

from array import array
xrange = getattr(__builtins__, 'xrange', range)
 
def version2():
    def to_rgb(pixels, buffer_len):
        for i in xrange(0, buffer_len - 2, 3):
            yield pixels[i + 2]
            yield pixels[i + 1]
            yield pixels[i]
 
    return array('H', to_rgb(pixels, buffer_len))

Mais mesurer révèle que cette solution est au final bien plus lente que celle de l’auteur :

import timeit
print(timeit.timeit('version1()', setup='from __main__ import version1'))
#16.9543120861
print(timeit.timeit('version2()', setup='from __main__ import version2'))
#42.7808477879

C’est bubulle qui va nous pondre une solution originale, élégante, mais surtout 5 fois plus rapide, ce qui est complètement inattendu :

def version3():
    pixels[2:buffer_len:3], pixels[0:buffer_len:3] = pixels[0:buffer_len:3], pixels[2:buffer_len:3]

Il utilise 3 astuces en une :

• Le swap de variable par unpacking: a, b = b, a. On sait que c’est plus rapide que deux assignation car Python saute la version intermédiaire.
• Le slicing avec pas (l[debut:fin:pas]), et il parcours donc son array de 3 en 3. Super malin ! Non seulement je n’y aurais jamais pensé, mais en prime j’aurais imaginé que créer deux nouveaux arrays à partir du premier serait lent.
• L’assignation de slices. Oui, on peut faire l[3:5] = [‘a’, ‘b’] en Python.

Le gain de perf est assez bluffant :

print(timeit.timeit('version3()', setup='from __main__ import version3'))
#1.84227705002

Pourquoi est-ce aussi rapide ? Je ne peux que faire des suppositions.

Je pense que faisant cette opération sans boucler explicitement, une bonne partie reste dans l’implémentation en C sous-jacente. On évite notamment pas mal d’appels à __getitem__ et __setitem__ qui sont toujours coûteux en Python.

Par curiosité, j’ai lancé le même test sur pypy :

0.359977960587
15.5387830734
0.750488996506

On voit l’effet de la compilation JIT puisque la boucle explicite redevient plus performante.

Mais du coup je me pose une question : pourquoi ma version, qui utilise une boucle explicite également, reste autant à la traine ?

Puis j’ai réalisé : l’array stock des entier C, mais quand on les sort, il faut les rewrapper dans des entiers Python. Donc ma boucle se tape cette conversion à chaque tour, et ça bouffe.

Je vire l’array :

def version4():
    def to_rgb(pixels, buffer_len):
        for i in xrange(0, buffer_len - 2, 3):
            yield pixels[i + 2]
            yield pixels[i + 1]
            yield pixels[i]
 
    return list(to_rgb(pixels, buffer_len))

Et pouf, le temps devient plus raisonnable.

Avec CPython : 1.84687900543
Avec Pypy : 0.274104118347

Sur CPython, la vitesse est du même ordre de grandeur que la version avec les slices.

Sur PyPy, la version est 2X plus rapide que les slices, est 25% plus rapide que la version originale.

Moralités :

• Les slices, c’est rapide. Et beau.
• Bien connaitre ses structures de données. Utiliser la mauvaise vous tue les perfs, et array sert à économiser de la place en mémoire, pas à gagner du CPU.
• Pypy, ça trace, mais le gain de perf peut se déplacer d’un algo à l’autre.