Функция замедления @jit

Я разрабатываю код оптимизации для сложной задачи эксплуатации пласта. Часть этого требует, чтобы я вычислил целевую функцию для большого количества потенциальных решений. Я тестирую оптимизатор на функции Розенброка и пытаюсь улучшить его скорость. Когда я профилировал код, я заметил, что вычисление целевой функции в цикле for было одним из узких мест кода, поэтому я разработал способ делать это параллельно для нескольких наборов переменных решения. У меня есть два калькулятора целевых функций: FO для одного набора переменных решения и P_FO для нескольких наборов переменных решения. Вычисление целевой функции — одна из самых медленных частей моего кода, поэтому я хотел бы еще больше ускорить процесс, используя @jit. Я протестировал обе функции с помощью @jit и обнаружил, что функция P_FO с @jit работает медленнее, чем без нее. Код ниже:

import time
import numpy as np
from numba import jit 

def FO(X):
    #Rosenbrock function
    ObjV=0
    for i in range(65-1):
        F=100*((X[i+1]-X[i]**2)+(X[i]-1)**2)
        ObjV+=F
    return ObjV

t0=time.time()
X=10+np.zeros(65)
for i in range(5000):
    FO(X)
t1 = time.time()
total = t1-t0
print("time FO="+str(total))

@jit
def FO(X):
    #Rosenbrock function
    ObjV=0
    for i in range(65-1):
        F=100*((X[i+1]-X[i]**2)+(X[i]-1)**2)
        ObjV+=F
    return ObjV

t0=time.time()
X=10+np.zeros(65)
for i in range(5000):
    FO(X)
t1 = time.time()
total = t1-t0
print("time FO with @jit="+str(total))



def P_FO(X):
    ObjV=np.zeros(X.shape[0])  
    for i in range(X.shape[1]-1):
        F=100*((X[:,i+1]-X[:,i]**2)+(X[:,i]-1)**2)
        ObjV+=F
    return ObjV       

t0=time.time()
X=10+np.zeros((65, 5000))
P_FO(X)
t1 = time.time()
total = t1-t0
print("time P_FO="+str(total))


@jit
def P_FO(X):
    ObjV=np.zeros(X.shape[0])  
    for i in range(X.shape[1]-1):
        F=100*((X[:,i+1]-X[:,i]**2)+(X[:,i]-1)**2)
        ObjV+=F
    return ObjV       

t0=time.time()
X=10+np.zeros((65, 5000))
P_FO(X)
t1 = time.time()
total = t1-t0
print("time P_FO with @jit="+str(total))

Результаты были:

time FO=0.523999929428
time FO with @jit=0.0720000267029
time P_FO=0.0380001068115
time P_FO with @jit=0.229000091553

Может ли кто-нибудь указать мне причину, по которой @jit замедляет параллельную целевую функцию P_FO? Это из-за использования np.zeros или, возможно, array.shape()?


person Kingle    schedule 20.07.2017    source источник


Ответы (1)


numba-функции компилируются лениво, т. е. до тех пор, пока не будут вызваны в первый раз, поэтому ваши тайминги отражают единовременные накладные расходы на компиляцию. Если я вызову каждую функцию один раз перед запуском секций синхронизации, я получу:

time FO=0.4103426933288574
time FO with @jit=0.0020008087158203125
time P_FO=0.04154801368713379
time P_FO with @jit=0.004002809524536133
person chrisb    schedule 21.07.2017
comment
Это имеет смысл. Поскольку я буду запускать его до 300 раз за симуляцию (и миллионы симуляций), это все равно может улучшить результаты. - person Kingle; 21.07.2017
comment
Просто любопытно - если функция @jit находится в другом файле .py, должна ли она компилироваться каждый раз при вызове? Или его компиляция один раз из моей основной функции ускорит его для остальных вызовов функций? - person Kingle; 21.07.2017
comment
Нет, после импорта компиляцию нужно выполнить только один раз. Более того, есть опция cache=True, которая может кэшировать компиляцию между сеансами Python. См. документы по адресу numba.pydata.org/numba-doc/latest. /user/jit.html#cache - person chrisb; 21.07.2017