我有一些使用 numpy 计算函数梯度的 python 代码，这是我的应用程序中的一个大瓶颈。所以，我最初的尝试是尝试用它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]))
除了最后 2 个参数之外的所有内容都是 numpy 数组，并且如 cython 函数签名中所述。python代码几乎相同，如果你愿意，我可以发布它，但它基本上是一样的。
我用 as 编译了整个东西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 ，据我所知，我做了它推荐的几乎所有事情。任何关于我下一步可以尝试的建议将不胜感激！