本文探讨了在Python中高效计算两组向量间稀疏欧氏距离的策略。针对传统方法中计算大量不必要距离的性能瓶颈，我们提出并实现了一种结合Numba加速和SciPy稀疏矩阵（CSR格式）的解决方案。该方法通过显式循环和条件判断，仅计算所需距离，并直接构建稀疏矩阵，显著提升了计算速度和内存效率，特别适用于大规模、高稀疏度的场景。
问题背景与传统方法的局限性
在许多数据分析和机器学习任务中，我们可能需要计算两组向量集合a和b之间的所有或部分欧氏距离。当需要计算的距离对数量远小于总的可能距离对数量时（即距离矩阵是稀疏的），传统的计算方法效率低下。
考虑以下场景：给定两个向量集合 A (形状为 (N, D)) 和 B (形状为 (M, D))，以及一个布尔掩码 M (形状为 (N, M))，其中 M[i, j] 为 True 表示需要计算 A[i] 和 B[j] 之间的距离。
一种常见的初始实现方式是计算所有 N * M 对距离，然后通过掩码 M 筛选出所需的部分：
import numpy as npA = np.array([[1, 2], [2, 3], [3, 4]]) # (3, 2)B = np.array([[4, 5], [5, 6], [6, 7], [7, 8], [8, 9]]) # (5, 2)M = np.array([[0, 0, 0, 1, 0], [1, 1, 0, 0, 0], [0, 0, 0, 0, 1]]) # (3, 5)# 计算所有差值矩阵diff = A[:,None] - B[None,:] # (3, 5, 2)# 计算所有欧氏距离distances = np.linalg.norm(diff, ord=2, axis=2) # (3, 5)# 应用掩码，将不需要的距离置零masked_distances = distances * M # (3, 5)print("原始方法计算的距离矩阵：n", masked_distances)
这种方法的问题在于，即使 M 中绝大多数元素为 False (即稀疏度很高)，它仍然计算了所有 N * M 对距离，导致大量的冗余计算和内存消耗，尤其当 N 和 M 达到数千甚至更高时，性能瓶颈将非常明显。尝试使用 np.vectorize 或稀疏数组处理高维 diff 矩阵也往往无法有效解决性能问题。
基于Numba和CSR稀疏矩阵的高效解决方案
为了解决上述性能问题，我们可以采用一种结合Numba即时编译和SciPy稀疏矩阵（特别是CSR格式）的方法。核心思想是：
立即学习“Python免费学习笔记（深入）”；
避免不必要的计算： 仅在掩码 M 对应位置为 True 时才计算欧氏距离。高效存储： 直接将计算出的非零距离构建成稀疏矩阵，避免存储大量的零值。性能优化： 利用Numba加速Python循环，使其性能接近C语言。
1. 导入所需库
import numba as nbimport numpy as npimport scipy.sparseimport math
2. Numba加速的欧氏距离函数
np.linalg.norm 在循环内部调用时会有一定的开销。为了最大限度地提高性能，我们可以编写一个Numba加速的自定义欧氏距离函数。
@nb.njit()def euclidean_distance(vec_a, vec_b): """ 计算两个向量之间的欧氏距离。 使用Numba进行JIT编译以提高性能。 """ acc = 0.0 for i in range(vec_a.shape[0]): acc += (vec_a[i] - vec_b[i]) ** 2 return math.sqrt(acc)
这个函数使用一个简单的循环计算平方差之和，然后取平方根，其结果与 np.linalg.norm(vec_a – vec_b) 相同，但在Numba的优化下，执行速度更快。
3. 核心计算与CSR矩阵构建逻辑
构建CSR（Compressed Sparse Row）矩阵需要三个核心数组：
data：存储非零元素的值。indices：存储每个非零元素所在的列索引。indptr：存储每行非零元素在 data 和 indices 数组中的起始位置。
以下Numba加速的函数负责遍历掩码，计算所需距离，并填充这些CSR数组。
@nb.njit()def masked_distance_inner(data, indices, indptr, matrix_a, matrix_b, mask): """ Numba加速的核心函数，根据掩码计算距离并填充CSR矩阵的data、indices、indptr数组。 """ write_pos = 0 # 当前写入data和indices数组的位置 N, M = matrix_a.shape[0], matrix_b.shape[0] for i in range(N): for j in range(M): if mask[i, j]: # 如果掩码为True，则计算距离并记录 data[write_pos] = euclidean_distance(matrix_a[i], matrix_b[j]) indices[write_pos] = j # 记录列索引 write_pos += 1 # 每行结束后，记录该行在data和indices数组中的结束位置（即下一行的起始位置） indptr[i + 1] = write_pos # 确保所有预分配的内存都被使用 assert write_pos == data.shape[0] assert write_pos == indices.shape[0] # data、indices、indptr数组通过参数修改，无需返回
这个函数是性能关键部分，它通过两个嵌套循环遍历所有可能的 (i, j) 对。只有当 mask[i, j] 为 True 时，才会调用 euclidean_distance 计算距离，并将结果存储到 data 数组中，同时记录其对应的列索引到 indices 数组。indptr 数组则负责标记每行非零元素的起始和结束位置，这是CSR格式的关键。
4. 封装函数：设置与CSR矩阵生成
为了方便使用，我们创建一个封装函数 masked_distance，它负责准备数据结构、调用Numba核心函数，并最终返回一个 scipy.sparse.csr_matrix 对象。
def masked_distance(matrix_a, matrix_b, mask): """ 计算两组向量间由掩码指定的稀疏欧氏距离，并返回一个CSR稀疏矩阵。 """ N, M = matrix_a.shape[0], matrix_b.shape[0] assert mask.shape == (N, M), "掩码形状必须与A和B的行数匹配。" # 将掩码转换为布尔类型，确保正确计数 mask = mask != 0 # 确定稀疏矩阵中非零元素的总数，用于预分配内存 sparse_length = mask.sum() # 预分配CSR矩阵所需的数组。无需初始化，Numba函数会直接写入。 data = np.empty(sparse_length, dtype='float64') # 存储距离值 indices = np.empty(sparse_length, dtype='int64') # 存储列索引 indptr = np.zeros(N + 1, dtype='int64') # 存储每行的起始索引 # 调用Numba加速的核心函数进行计算和填充 masked_distance_inner(data, indices, indptr, matrix_a, matrix_b, mask) # 使用填充好的数组构建scipy.sparse.csr_matrix return scipy.sparse.csr_matrix((data, indices, indptr), shape=(N, M))
这个封装函数首先进行输入校验，然后计算稀疏矩阵中非零元素的总数 sparse_length。这是非常关键的一步，因为它允许我们精确地预分配 data 和 indices 数组的内存，避免了动态调整大小的开销。indptr 数组则需要预先用零初始化，并在 masked_distance_inner 中逐步填充。
示例与性能分析
为了演示和验证性能，我们创建大型随机数据集：
# 创建大型随机数据集进行测试A_big = np.random.rand(2000, 10) # 2000个10维向量B_big = np.random.rand(4000, 10) # 4000个10维向量# 创建一个非常稀疏的掩码（0.1%的元素为True）M_big = np.random.rand(A_big.shape[0], B_big.shape[0]) < 0.001# 运行基准测试# %timeit masked_distance(A_big, B_big, M_big)
在实际测试中，对于 A_big (2000, 10) 和 B_big (4000, 10)，M_big 稀疏度为0.1%的场景，这种方法比原始的“计算所有再掩码”的方法快了约40倍。当向量维度更高（例如1000维）时，加速比甚至可以达到1000倍。
原始方法（计算所有距离再掩码）的性能示例：556 ms ± 3.74 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Numba加速和CSR构建方法的性能示例：13.5 ms ± 66.6 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
这种显著的性能提升主要归因于：
避免冗余计算： 仅计算掩码指定的距离，大大减少了浮点运算量。Numba加速： 将Python循环转换为高效的机器码，消除了Python解释器的开销。内存效率： CSR矩阵只存储非零元素及其位置信息，节省了大量内存。
注意事项与优化建议
Numba的JIT编译： euclidean_distance 和 masked_distance_inner 函数上的 @nb.njit() 装饰器是性能提升的关键。它指示Numba在函数首次调用时将其编译为优化的机器码。自定义欧氏距离： 实践证明，在Numba环境中，自定义的循环计算欧氏距离通常比调用 np.linalg.norm 更快，因为它避免了函数调用的开销和NumPy通用函数的额外检查。数据类型优化：距离值 (data)： 如果对精度要求不高，可以将 float64 替换为 float32，这可以减少内存占用并可能略微提高计算速度。索引 (indices, indptr)： 如果矩阵的维度（行数或列数）以及非零元素的总数小于 231，可以将 int64 替换为 int32，进一步节省内存。稀疏度影响： 这种方法的性能优势随着掩码稀疏度的增加而更加明显。如果掩码非常稠密（即需要计算大部分距离），那么原始的NumPy向量化方法可能仍然具有竞争力，因为Numba循环的固定开销可能会抵消部分优势。正确性验证： 在实际应用中，务必使用 np.allclose() 等方法验证新算法的计算结果与原始正确结果的一致性。
总结
通过结合Numba的即时编译能力和SciPy的CSR稀疏矩阵结构，我们提供了一种在Python中高效计算稀疏欧氏距离的专业教程方案。该方法通过有条件地计算距离并直接构建稀疏数据结构，有效解决了传统方法中存在的性能瓶颈和内存浪费问题。尤其在处理大规模、高稀疏度的向量集合时，这种优化策略能够带来显著的性能提升，是处理此类问题的理想选择。
以上就是优化Python中稀疏向量对欧氏距离计算的性能的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/1374654.html
微信扫一扫 
支付宝扫一扫 
