Проблема с производительностью при автоматическом распараллеливании в `pytorch` ⇐ Python

[Anonymous]

У меня проблемы с кодом Python, использующим pytorch. Детали немного сложны (код является частью квантово-механического расчета), но структура кода очень проста и выглядит примерно так:

Код: Выделить всё

# 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)

То же самое вычисление выполняется, скажем, для 100000 наборов импульсов. Распараллеливание этого с помощью Fortran потребует добавления простого параллельного выполнения !$omp.
Я использую Python частично как возможность обучения, а частично потому, что позже мне захочется автоматически рассчитывать градиенты по отношению к некоторым параметрам. К сожалению, при измерении производительности кода я получаю следующую зависимость между временем выполнения и количеством используемых ядер:


Поскольку памяти много и вычисления можно легко распараллелить по первому индексу p, я ожидаю, что вместо этого эта зависимость будет гораздо сильнее уменьшаться. Например, при 20 потоках я ожидаю, что время выполнения составит около 7 секунд: 1/20 времени одного потока.
Я предполагаю, что автоматическое распараллеливание , не является оптимальным. Можно ли явно указать, что расчет необходимо выполнять параллельно для каждого i в p[i , : , :]? Я мог бы запустить 20 (или более) потоков Python и при каждом запуске с потоками == interop threads == 1, но я надеюсь, что есть более простое и элегантное решение.
Советы и комментарии будут очень признательны.
PS Для одного потока производительность аналогична Fortran.

Подробнее здесь: https://stackoverflow.com/questions/798 ... in-pytorch