numpy 向量化一个函数会减慢它的速度吗？

Question

我正在编写一个程序，在该程序中我试图查看给定的红移如何很好地获得在光谱中检测到的一组线以匹配原子线数据库。红移越接近线条重叠，“分数”越低，红移正确的机会就越高。

我通过遍历一系列可能的红移，计算每个红移的分数来做到这一点。在那个外循环中，我在检测到的行集合中的每一行中循环以计算它的 sub_score，然后对该内循环求和以获得总分。

我尝试用 numpy 向量化内部循环，但令人惊讶的是它实际上减慢了执行速度。在给出的示例中，嵌套 for 循环在我的笔记本电脑上执行需要约 2.6 秒，而内部带有 numpy 的单个 for 循环需要约 5.3 秒。

为什么矢量化内部循环会减慢速度？有没有我想念的更好的方法？

import numpy as np
import time

def find_nearest_line(lines, energies):
    # Return the indices (from the lines vector) that are closest to the energies vector supplied
    # Vectorized with help from 
    energies = np.expand_dims(energies, axis=-1)
    idx = np.abs(lines / energies - 1).argmin(axis=-1)
    return idx

def calculate_distance_to_line(lines, energies):
    # Returns distance between an array of lines and an array of energies)
    z = (lines / energies) - 1
    return z

rng = np.random.default_rng(2021)
det_lines = rng.random(1000)
atom_lines = rng.random(10000)

redshifts = np.linspace(-0.1, 0.1, 100)

# loop version
start=time.time()
scores = []
for redshift in redshifts:
    atom_lines_shifted = atom_lines * (1 + redshift)
    score = 0
    for det_line in det_lines:
        closest_atom_line = find_nearest_line(atom_lines_shifted, det_line)
        score += abs(calculate_distance_to_line(atom_lines_shifted[closest_atom_line], det_line))
    scores.append(score)
print(time.time()-start)
print(scores)

# (semi)-vectorized version
start=time.time()
scores = []
for redshift in redshifts:
    atom_lines_shifted = atom_lines * (1 + redshift)
    closest_atom_lines = find_nearest_line(atom_lines_shifted, det_lines)
    score = np.sum(np.abs(calculate_distance_to_line(atom_lines_shifted[closest_atom_lines], det_lines)))
    scores.append(score)
print(time.time()-start)
print(scores)

Answer 1

Numpy 代码通常会创建许多临时数组。例如，您的函数 find_nearest_line 就是这种情况。同时处理 det_lines 的所有项目将导致创建许多相对较大的数组（每个数组 1000 * 10_000 * 8 = 76 MiB）。问题是 大数组通常不适合 CPU 缓存 。如果是这样，数组需要存储在 RAM 中，吞吐量 低得多 并且延迟 延迟高得多 。此外， and results often in more page faults（由于大多数默认标准分配器的实际实现）。有时使用大数组会更快，因为 CPython 解释器的开销很大，但是这两种策略在实践中都是低效的。

问题是算法效率不高。事实上，您可以 对数组 排序并使用 二分搜索 来更有效地找到最接近的值。 np.searchsorted 完成大部分工作，但它仅 returns 大于（或等于）目标值的最接近值的索引。因此，需要执行一些额外的操作才能获得最接近的值（可能大于或小于目标值）。请注意，由于二进制搜索，此算法不会生成巨大的数组。

scores = []
n = atom_lines.size
m = det_lines.size
line_idx = np.arange(m)

for redshift in redshifts:
    atom_lines_shifted = atom_lines * (1 + redshift)
    sorted_atom_lines_shifted = np.sort(atom_lines_shifted)
    close_idx = np.searchsorted(sorted_atom_lines_shifted, det_lines)
    lower_bound = sorted_atom_lines_shifted[np.maximum(close_idx - 1, 0)]
    upper_bound = sorted_atom_lines_shifted[np.minimum(close_idx, n - 1)]
    bounds = np.hstack((lower_bound[:, None], upper_bound[:, None]))
    closest_bound_idx = find_nearest_line(bounds, det_lines)
    close_values = bounds[line_idx, closest_bound_idx]
    score = np.sum(np.abs(calculate_distance_to_line(close_values, det_lines)))
    scores.append(score)

由于atom_lines没有修改，乘法保留顺序，算法可以通过直接排序atom_lines进一步优化：

scores = []
n = atom_lines.size
m = det_lines.size
line_idx = np.arange(m)
sorted_atom_lines = np.sort(atom_lines)

for redshift in redshifts:
    sorted_atom_lines_shifted = sorted_atom_lines * (1 + redshift)
    close_idx = np.searchsorted(sorted_atom_lines_shifted, det_lines)
    lower_bound = sorted_atom_lines_shifted[np.maximum(close_idx - 1, 0)]
    upper_bound = sorted_atom_lines_shifted[np.minimum(close_idx, n - 1)]
    bounds = np.hstack((lower_bound[:, None], upper_bound[:, None]))
    closest_bound_idx = find_nearest_line(bounds, det_lines)
    close_values = bounds[line_idx, closest_bound_idx]
    score = np.sum(np.abs(calculate_distance_to_line(close_values, det_lines)))
    scores.append(score)

最后一个实现在我的机器上快 300 倍。

numpy 向量化一个函数会减慢它的速度吗？

numpy vectorizing a function slows it down?

python

performance

numpy