Python: Primitive Summe eines 3D-Arrays vs. Numba JIT Benchmark

Diese Funktion summiert alle Werte des gegebenen 3D-Arrays:

def primitive_pixel_sum(frame):
    result = 0.0
    for x in range(frame.shape[0]):
        for y in range(frame.shape[1]):
            for z in range(frame.shape[2]):
                result += frame[x,y,z]
    return result

Während die folgende Funktion exakt denselben Algorithmus verwendet, aber mit numba.jit:

import numba

@numba.jit
def numba_pixel_sum(frame):
    result = 0.0
    for x in range(frame.shape[0]):
        for y in range(frame.shape[1]):
            for z in range(frame.shape[2]):
                result += frame[x,y,z]
    return result

Wir können sie in Jupyter benchmarken mit

%%timeit
primitive_pixel_sum(frame)

und

%%timeit
numba_pixel_sum(frame)

Ergebnisse

Wir haben dies mit einem zufälligen Kamerabild aus OpenCV der Shape (480, 640, 3) getestet

primitive_pixel_sum():

1.78 s ± 253 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

numba_pixel_sum():

4.06 ms ± 413 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Aus diesen Ergebnissen sollte klar hervorgehen, dass die Numba-Version 438-mal schneller ist als die primitive Version.

Beachte, dass beim Kompilieren komplexer Funktionen mit numba.jit viele Millisekunden oder sogar Sekunden vergehen können — möglicherweise länger, als eine einfache Python-Funktion benötigen würde.

Da Numba so einfach zu verwenden ist, ist meine Empfehlung, es einfach für jede Funktion auszuprobieren, bei der du vermutest, dass sie viel CPU-Zeit verbraucht. Mit der Zeit wirst du ein Gespür dafür entwickeln, bei welchen Funktionen sich Numba lohnt und bei welchen es gar nicht funktioniert oder insgesamt langsamer ist als reines Python.

Denke daran, dass du oft auch NumPy-Funktionen verwenden kannst, um dasselbe Ergebnis zu erzielen. In unserem Beispiel könntest du dasselbe erreichen mit

np.sum(frame)

was sogar noch schneller ist als Numba:

%%timeit
np.sum(frame)

Ergebnis:

2.5 ms ± 7.17 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)