Wie beschleunigt man dict finden Aufgabe in Numba?

Question

Ich habe eine Python-Funktion und möchte sie mit Numba beschleunigen. Der zeitaufwendigste Teil ist eine Suche nach Tupeln in einem Set/Dict. Kann jemand eine Idee geben, wie man das löst?

idNeg = np.array([[1,2,3], [4,5,6], ..., [7,8,9]]) 
validSet = {(1,2,3):True, (5,4,3):True, ..., (2,5,3):True} 
@jit 
def CalcNeg(idNeg, validSet): 
    l = len(idNeg) 
    for j in xrange(l): 
     corplc = np.random.choice([0, 2]) 
     idNeg[j, corplc] = random.randrange(0, VE.shape[0]) 
     while validSet.has_key((idNeg[j, 0], idNeg[j, 1], idNeg[j, 2])): 
      idNeg[j, corplc] = random.randrange(0, VE.shape[0]) 
    return idNeg 

ich so versucht, aber die Geschwindigkeit nicht den Code ohne @jit ändern kann.

Answer 1

Ich bin nicht 100% klar über die genaue Art der gewünschten Eingaben (Inhalt oder typische Formen), aber der Schlüssel zu guten Leistungszuwächsen in Numba ist in der Lage zu jit eine Funktion in nopython Modus (im Gegensatz zu Python-Objektmodus). Die ursprüngliche Funktion verwendete Datenstrukturen, insbesondere ein Diktat, das momentan nicht unterstützt wird.

Noch einmal, ich kenne nicht den genauen Anwendungsfall oder ob die folgende Änderung gültig ist, aber ich nahm die validSet dict und konvertierte die Schlüssel zu einem realen Set-Objekt, wo der Wert im Schlüssel, Wert-Paar war True.

Als Beispiel:

import numpy as np 
import numba as nb 
import random 

# Original function 
def CalcNeg(idNeg, validSet, N): 
    l = len(idNeg) 
    for j in xrange(l): 
     corplc = np.random.choice([0, 2]) 
     idNeg[j, corplc] = random.randrange(0, N) 
     while validSet.has_key((idNeg[j, 0], idNeg[j, 1], idNeg[j, 2])): 
      idNeg[j, corplc] = random.randrange(0, N) 
    return idNeg 

# Modified version, compiled in nopython mode (njit) 
@nb.njit 
def CalcNeg2(idNeg, validSet, N): 
    l = len(idNeg) 
    c = np.array([0,2]) 
    for j in xrange(l): 
     corplc = np.random.choice(c) 
     idNeg[j, corplc] = random.randrange(0, N) 
     #while validSet.has_key((idNeg[j, 0], idNeg[j, 1], idNeg[j, 2])): 
     while (idNeg[j, 0], idNeg[j, 1], idNeg[j, 2]) in validSet: 
      idNeg[j, corplc] = random.randrange(0, N) 
    return idNeg 

# Some test data 
N = 40 
M = 2000 
idNeg = np.random.random_integers(0, N, size=(M,3)) 
tmp = np.random.random_integers(0, N, size=(M,3)) 
validSet = {tuple(tmp[k,:]): True for k in xrange(tmp.shape[0])} 

# convert validSet to real python set for keys with value == True 
_validSet = {k for k,v in validSet.iteritems() if v is True} 

Und jetzt einige Timings von einem ipython Notebook mit der %timeit Magie:

%timeit CalcNeg(idNeg, validSet, N) 
100 loops, best of 3: 7.84 ms per loop 

%timeit CalcNeg2(idNeg, _validSet, N) 
1000 loops, best of 3: 444 µs per loop 

So auf meiner Maschine, die ein 18-facher Geschwindigkeit-up ist. Ich benutze Numba 0.25. Nur als Hinweis, der Wechsel zu einem Set in der ursprünglichen Python-Funktion gibt einen kleinen Unterschied, aber mehr wie eine 25% ige Beschleunigung.

Wenn die Testdaten unrealistisch sind, oder wenn die Umwandlung des Diktats in eine Menge nicht angemessen ist, lassen Sie es mich wissen. Ohne genauere Angaben ist es schwer zu sagen, wie man das Problem angeht.