У меня есть некоторый код на python, который использует numpy, который вычисляет градиент функции, и это большое узкое место в моем приложении. Итак, моей первоначальной попыткой было попытаться использовать Cython для улучшения производительности.
Итак, используя онлайн-руководства, я смог легко перенести это на Cython, но получил очень умеренное ускорение примерно на 15%. Функция содержит много циклов, и я надеялся, что Cython даст гораздо лучшее улучшение.
Код Cython выглядит следующим образом. Ниже приведены вспомогательные функции, которые вызываются только из Cython.
cimport numpy as np
cimport cython
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
cdef cget_cubic_bspline_weight(double u):
u = fabs(u)
if u < 2.0:
if u < 1.0:
return 2.0 / 3.0 - u ** 2 + 0.5 * u ** 3
else:
return ((2.0 - u) ** 3) / 6.0
return 0.0
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
cdef cget_cubic_spline_first_der_weight(double u):
cdef double o = u
u = fabs(u)
cdef double v
if u < 2.0:
if u < 1.0:
return (1.5 * u - 2.0) * o
else:
u -= 2.0
v = -0.5 * u * u
if o < 0.0:
return -v
return v
return 0.0;
Ниже приводится основная функция, которая вычисляет градиент.
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
cpdef gradient_2d(np.ndarray[double, ndim=2, mode="c"] reference,
np.ndarray[double, ndim=2, mode="c"] warped,
np.ndarray[double, ndim=5, mode="fortran"] warped_gradient,
np.ndarray[double, ndim=5, mode="fortran"] result_gradient,
double[:] entropies,
np.ndarray[double, ndim=2, mode="c"] jhlog,
np.ndarray[double, ndim=2, mode="fortran"] reflog,
np.ndarray[double, ndim=2, mode="fortran"] warlog,
int[:] bins,
int height, int width):
war_x = warped_gradient[..., 0]
war_y = warped_gradient[..., 1]
res_x = result_gradient[..., 0]
res_y = result_gradient[..., 1]
nmi = (entropies[0] + entropies[1]) / entropies[2]
for y in range(height):
for x in range(width):
ref = reference[x, y]
war = warped[x, y]
jd = [0.0] * 2
rd = [0.0] * 2
wd = [0.0] * 2
for r in range(int(ref - 1.0), int(ref + 3.0)):
if (-1 < r and r < bins[0]):
for w in range(int(war - 1.0), int(war + 3.0)):
if (-1 < w and w < bins[1]):
c = cget_cubic_bspline_weight(ref - float(r)) * \
cget_cubic_spline_first_der_weight(war - float(w))
jl = jhlog[r, w]
rl = reflog[r, 0]
wl = warlog[0, w]
jd[0] += c * war_x[x, y] * jl
rd[0] += c * war_x[x, y] * rl
wd[0] += c * war_x[x, y] * wl
jd[1] += c * war_y[x, y] * jl
rd[1] += c * war_y[x, y] * rl
wd[1] += c * war_y[x, y] * wl
res_x[x, y] = (rd[0] + wd[0] - nmi * jd[0]) / (entropies[2] * entropies[3])
res_y[x, y] = (rd[1] + wd[1] - nmi * jd[1]) / (entropies[2] * entropies[3])
Теперь я называю это как:
speed.gradient_2d(self.rdata, self.wdata, warped_grad_image,
result_gradient.data, self.entropies,
self.jhlog, self.reflog, self.warlog, self.bins,
int(self.rdata.shape[1]), int(self.rdata.shape[0]))
Все, кроме двух последних параметров, являются массивами-пустышками и соответствуют описанию в сигнатуре функции cython. Код на Python почти такой же, и я могу опубликовать его, если хотите, но в основном он такой же.
Я скомпилировал все это с setup.py как:
from distutils.core import setup
from distutils.extension import Extension
from Cython.Build import cythonize
import numpy
ext = Extension("speed",
sources=["perf/speed.pyx"],
include_dirs=[numpy.get_include()],
language="c++",
libraries=[],
extra_link_args=[])
setup(ext_modules = cythonize([ext]))
Опять же, поскольку в моем коде так много циклов, у меня сложилось впечатление, что версия Cython будет намного быстрее, но я получу улучшение только на 15%. Я следовал этому руководству для реализации: http://docs.cython.org/en/latest/src/userguide/numpy_tutorial.html и, насколько я могу судить, я сделал почти все, что он рекомендует. Буду очень признателен за любые предложения о том, что я мог бы попробовать дальше!