Leistungsproblem bei der automatischen Parallelisierung in „Pytorch“. ⇐ Python

[Anonymous]

Ich habe Probleme mit Python-Code, der Pytorch verwendet. Die Details sind etwas kompliziert (der Code ist Teil einer quantenmechanischen Berechnung), aber die Codestruktur ist sehr einfach und sieht mehr oder weniger so aus:

Code: Select all

# p is a batch containing 100000 sets of momenta.
# Each set contains four vectors in 3 dimensions.
p = momenta[startbatch : endbatch]
# p.shape : (100000 , 4 , 3)

# It should be easy to parallelize the following
# with respect to the first index of `p`:
result = 1.0 # * 
# result.shape : (100000 , 16 , 16)

Die gleiche Berechnung wird beispielsweise für 100.000 Impulssätze durchgeführt. Um dies mit Fortran zu parallelisieren, müsste ein einfaches !$omp-Parallel-Do hinzugefügt werden.
Ich verwende Python teilweise als Lernmöglichkeit und teilweise, weil ich später Farbverläufe in Bezug auf einige Parameter automatisch berechnen möchte. Leider erhalte ich beim Messen der Leistung des Codes die folgende Beziehung zwischen der Ausführungszeit und der Anzahl der verwendeten Kerne:


Da Speicher reichlich vorhanden ist und die Berechnung leicht entlang des ersten Index von p parallelisiert werden kann würde ich erwarten, dass diese Beziehung stattdessen viel stärker abnehmen würde. Beispielsweise würde ich bei 20 Threads eine Ausführungszeit von etwa 7 Sekunden erwarten: 1/20 der Zeit eines einzelnen Threads.
Ich vermute, dass die automatische Parallelisierung von nicht optimal ist. Kann explizit angegeben werden, dass die Berechnung für jedes i in p[i , : , :] parallel durchgeführt werden muss? Ich könnte 20 (oder mehr) Python-Threads ausführen und in jedem with threads == interop threads == 1 ausführen, aber ich hoffe, dass es eine einfachere/elegantere Lösung gibt.
Tipps und Kommentare wären sehr dankbar.
PS: Für einen einzelnen Thread ist die Leistung ähnlich wie bei Fortran.