我写了一个这样的卷积函数:
import numpy as np
import numba as nb
# Generate sample input data
num_chans = 111
num_bins = 47998
num_rad = 8
num_col = 1000
rng = np.random.default_rng()
wvl_sensor = rng.uniform(low=1000, high=11000, size=(num_chans, num_col))
fwhm_sensor = rng.uniform(low=0.01, high=2.0, size=num_chans)
wvl_lut = rng.uniform(low=1000, high=11000, size=num_bins)
rad_lut = rng.uniform(low=0, high=1, size=(num_rad, num_bins))
# Original convolution implementation
def original_convolve(wvl_sensor, fwhm_sensor, wvl_lut, rad_lut):
sigma = fwhm_sensor / (2.0 * np.sqrt(2.0 * np.log(2.0)))
var = sigma ** 2
denom = (2 * np.pi * var) ** 0.5
numer = np.exp(-(wvl_lut[:, None] - wvl_sensor[None, :])**2 / (2*var))
response = numer / denom
response /= response.sum(axis=0)
resampled = np.dot(rad_lut, response)
return resampled
字符串
numpy版本运行时间约为45秒:
# numpy version
num_chans, num_col = wvl_sensor.shape
num_bins = wvl_lut.shape[0]
num_rad = rad_lut.shape[0]
original_res = np.empty((num_col, num_rad, num_chans), dtype=np.float64)
for x in range(wvl_sensor.shape[1]):
original_res[x, :, :] = original_convolve(wvl_sensor[:, x], fwhm_sensor, wvl_lut, rad_lut)
型
我试着用numba加速它:
@nb.jit(nopython=True)
def numba_convolve(wvl_sensor, fwhm_sensor, wvl_lut, rad_lut):
num_chans, num_col = wvl_sensor.shape
num_bins = wvl_lut.shape[0]
num_rad = rad_lut.shape[0]
output = np.empty((num_col, num_rad, num_chans), dtype=np.float64)
sigma = fwhm_sensor / (2.0 * np.sqrt(2.0 * np.log(2.0)))
var = sigma ** 2
denom = (2 * np.pi * var) ** 0.5
for x in nb.prange(num_col):
numer = np.exp(-(wvl_lut[:, None] - wvl_sensor[None, :, x])**2 / (2*var))
response = numer / denom
response /= response.sum(axis=0)
resampled = np.dot(rad_lut, response)
output[x, :, :] = resampled
return output
型
它仍然花费大约32秒。请注意,如果我使用@nb.jit(nopython=True, parallel=True)
,输出是全零值。
有什么想法可以正确应用numba?或者改进卷积函数?
1条答案
按热度按时间t5fffqht1#
您可以将Numba实现中类似Numpy的代码替换为普通循环。Numba喜欢普通循环。这种策略有助于避免创建许多小的临时数组或一些大的临时数组。这两者都会导致可扩展性问题。然后您可以并行计算外部循环(就像您所做的那样)。
这种策略还有其他好处:
np.dot
使用一个 *BLAS库 *,它通常已经是并行的,所以在并行代码中使用它可能会导致性能问题(尽管它可以在应用程序级别进行调优);下面是生成的代码:
字符串
输出结果在我的机器上是正确的(有一个~ 1 e-16的错误)。注意,初始实现产生了像这样的NaN值。
结果
以下是我的i5- 9600 KF CPU(6核)的性能结果:
型
因此,这个新的并行Numba实现比最初的实现快5.74倍。它也比使用Numexpr快。这在我的CPU上接近最佳,使用默认的指数函数。
注意事项
在x86-64 CPU上使用英特尔SVML库的指数函数,代码肯定会更快一些。
如果你有一个服务器端的GPU,那么代码可以更有效地计算(因为这种计算是GPU友好的)。主流PC GPU可能不会比CPU更快地计算,因为主流PC GPU是为简单精度而设计的。
有了simple-precision,代码在CPU和GPU上都可以明显更快,特别是如果你不需要非常高的精度。这是因为: