numpy 如何加速卷积函数?

q3aa0525  于 2024-01-08  发布在  其他
关注(0)|答案(1)|浏览(218)

我写了一个这样的卷积函数:

  1. import numpy as np
  2. import numba as nb
  3. # Generate sample input data
  4. num_chans = 111
  5. num_bins = 47998
  6. num_rad = 8
  7. num_col = 1000
  8. rng = np.random.default_rng()
  9. wvl_sensor = rng.uniform(low=1000, high=11000, size=(num_chans, num_col))
  10. fwhm_sensor = rng.uniform(low=0.01, high=2.0, size=num_chans)
  11. wvl_lut = rng.uniform(low=1000, high=11000, size=num_bins)
  12. rad_lut = rng.uniform(low=0, high=1, size=(num_rad, num_bins))
  13. # Original convolution implementation
  14. def original_convolve(wvl_sensor, fwhm_sensor, wvl_lut, rad_lut):
  15. sigma = fwhm_sensor / (2.0 * np.sqrt(2.0 * np.log(2.0)))
  16. var = sigma ** 2
  17. denom = (2 * np.pi * var) ** 0.5
  18. numer = np.exp(-(wvl_lut[:, None] - wvl_sensor[None, :])**2 / (2*var))
  19. response = numer / denom
  20. response /= response.sum(axis=0)
  21. resampled = np.dot(rad_lut, response)
  22. return resampled

字符串
numpy版本运行时间约为45秒:

  1. # numpy version
  2. num_chans, num_col = wvl_sensor.shape
  3. num_bins = wvl_lut.shape[0]
  4. num_rad = rad_lut.shape[0]
  5. original_res = np.empty((num_col, num_rad, num_chans), dtype=np.float64)
  6. for x in range(wvl_sensor.shape[1]):
  7. original_res[x, :, :] = original_convolve(wvl_sensor[:, x], fwhm_sensor, wvl_lut, rad_lut)


我试着用numba加速它:

  1. @nb.jit(nopython=True)
  2. def numba_convolve(wvl_sensor, fwhm_sensor, wvl_lut, rad_lut):
  3. num_chans, num_col = wvl_sensor.shape
  4. num_bins = wvl_lut.shape[0]
  5. num_rad = rad_lut.shape[0]
  6. output = np.empty((num_col, num_rad, num_chans), dtype=np.float64)
  7. sigma = fwhm_sensor / (2.0 * np.sqrt(2.0 * np.log(2.0)))
  8. var = sigma ** 2
  9. denom = (2 * np.pi * var) ** 0.5
  10. for x in nb.prange(num_col):
  11. numer = np.exp(-(wvl_lut[:, None] - wvl_sensor[None, :, x])**2 / (2*var))
  12. response = numer / denom
  13. response /= response.sum(axis=0)
  14. resampled = np.dot(rad_lut, response)
  15. output[x, :, :] = resampled
  16. return output


它仍然花费大约32秒。请注意,如果我使用@nb.jit(nopython=True, parallel=True),输出是全零值。
有什么想法可以正确应用numba?或者改进卷积函数?

t5fffqht

t5fffqht1#

您可以将Numba实现中类似Numpy的代码替换为普通循环。Numba喜欢普通循环。这种策略有助于避免创建许多小的临时数组或一些大的临时数组。这两者都会导致可扩展性问题。然后您可以并行计算外部循环(就像您所做的那样)。
这种策略还有其他好处:

  • np.dot使用一个 *BLAS库 *,它通常已经是并行的,所以在并行代码中使用它可能会导致性能问题(尽管它可以在应用程序级别进行调优);
  • 并行Numba代码中的高级Numpy特性往往更 * 虚假 *(由于Numba本身),因此使用不太高级的特性可以避免此类问题(我遇到过几次这个问题)。

下面是生成的代码:

  1. @nb.jit(nopython=True, parallel=True)
  2. def numba_convolve_faster(wvl_sensor, fwhm_sensor, wvl_lut, rad_lut):
  3. num_chans, num_col = wvl_sensor.shape
  4. num_bins = wvl_lut.shape[0]
  5. num_rad = rad_lut.shape[0]
  6. original_res = np.empty((num_col, num_rad, num_chans), dtype=np.float64)
  7. sigma = fwhm_sensor / (2.0 * np.sqrt(2.0 * np.log(2.0)))
  8. var = sigma ** 2
  9. denom = (2 * np.pi * var) ** 0.5
  10. inv_denom = 1.0 / denom
  11. factor = -1 / (2*var)
  12. for x in nb.prange(wvl_sensor.shape[1]):
  13. wvl_sensor_col = wvl_sensor[:, x].copy()
  14. response = np.empty(num_bins)
  15. for j in range(num_chans):
  16. response_sum = 0.0
  17. for i in range(num_bins):
  18. diff = wvl_lut[i] - wvl_sensor_col[j]
  19. response[i] = np.exp(diff * diff * factor[j]) * inv_denom[j]
  20. response_sum += response[i]
  21. inv_response_sum = 1.0 / response_sum
  22. for i in range(num_bins):
  23. response[i] *= inv_response_sum
  24. for k in range(num_rad):
  25. s = 0.0
  26. for i in range(num_bins):
  27. s += rad_lut[k, i] * response[i]
  28. original_res[x, k, j] = s
  29. return original_res
  30. res = numba_convolve_faster(wvl_sensor, fwhm_sensor, wvl_lut, rad_lut)

字符串
输出结果在我的机器上是正确的(有一个~ 1 e-16的错误)。注意,初始实现产生了像这样的NaN值。

结果

以下是我的i5- 9600 KF CPU(6核)的性能结果:

  1. OP Numpy code: 4min 37s
  2. OP Numba code: 4min 24s
  3. With Numexpr: 1min 13s
  4. This Numba code: 46s <-----


因此,这个新的并行Numba实现比最初的实现快5.74倍。它也比使用Numexpr快。这在我的CPU上接近最佳,使用默认的指数函数。

注意事项

在x86-64 CPU上使用英特尔SVML库的指数函数,代码肯定会更快一些。
如果你有一个服务器端的GPU,那么代码可以更有效地计算(因为这种计算是GPU友好的)。主流PC GPU可能不会比CPU更快地计算,因为主流PC GPU是为简单精度而设计的。
有了simple-precision,代码在CPU和GPU上都可以明显更快,特别是如果你不需要非常高的精度。这是因为:

  • 简单精度的指数近似可以用一种非常SIMD友好的方式计算,而精确的双精度指数函数却不能(通常不值得);
  • 一些基本的简单精度运算比双精度运算更快(消耗更少的能量);
  • SIMD寄存器可以比双精度寄存器多保存两倍的简单精度数(通常每个SIMD操作的成本相同,因此吞吐量是简单精度数的两倍)。
展开查看全部

相关问题