sparse_dot_topn

sparse_dot_topn 提供了一种快速执行稀疏矩阵乘法并选择乘法结果前 n 个值的方法。

在实践中，比较非常大的特征向量并选取最佳匹配，通常需要执行稀疏矩阵乘法，然后选择乘法结果的前 n 个值。

sparse_dot_topn 提供了一种（并行化的）稀疏矩阵乘法实现，集成了选择前 n 个值的功能，显著降低了内存占用并提高了性能。在 Apple M2 Pro 上，对两个 20k x 193k 的 TF-IDF 矩阵进行操作时，sparse_dot_topn 在保留每行前 10 个值并利用 8 个核心的情况下，速度可以提高 6 倍。详细信息请查看 benchmark 目录。

使用方法

sp_matmul_topn 支持 {CSR, CSC, COO} 格式的矩阵，数据类型为 {32, 64}位 {整数, 浮点数}。请注意，COO 和 CSC 输入会被转换为 CSR 格式，因此速度较慢。两个可以进一步减少内存需求的选项是 threshold 和 density。可以选择对值进行排序，使得每行的第一列包含最大值。注意，sp_matmul_topn(A, B, top_n=B.shape[1]) 等同于 sp_matmul(A, B) 和 A.dot(B)。

如果你正在从 v0.* 版本迁移，请查看下面的迁移指南了解详情。

import scipy.sparse as sparse
from sparse_dot_topn import sp_matmul, sp_matmul_topn

A = sparse.random(1000, 100, density=0.1, format="csr")
B = sparse.random(100, 2000, density=0.1, format="csr")

# 计算 C 并保留每行前 10 个值
C = sp_matmul_topn(A, B, top_n=10)

# 或者使用并行化矩阵乘法，不选择前 n 个值
C = sp_matmul(A, B, n_threads=2)
# 或者选择前 n 个值
C = sp_matmul_topn(A, B, top_n=10, n_threads=2)

# 如果你只对高于某个阈值的值感兴趣
C = sp_matmul_topn(A, B, top_n=10, threshold=0.8)

# 如果设置了阈值，我们无法轻易预先确定非零项的数量。因此，我们为 `ceil(top_n * A.shap[0] * density)`
# 个非零项分配内存。你可以设置预期密度来减少预分配的项数。请注意，如果我们分配的内存不足，
# 将需要进行昂贵的复制操作。
C = sp_matmul_topn(A, B, top_n=10, threshold=0.8, density=0.1)

安装

sparse_dot_topn 为 CPython 3.8 到 3.12 提供以下平台的轮子：

Windows (64位)
Linux (64位)
MacOS (x86 和 ARM)

pip install sparse_dot_topn

sparse_dot_topn 依赖于 C++ 扩展来执行计算密集型的乘法例程。 注意，轮子中集成/附带了 OpenMP，以提供开箱即用的并行化功能。 这可能会在与其他集成 OpenMP 的库（如 PyTorch）一起使用时引起问题。 如果遇到任何 OpenMP 相关的问题，请查看 INSTALLATION.md 寻求帮助，或者在运行函数时不指定 n_threads 参数。

从源代码安装需要兼容 C++17 的编译器。如果你有可用的编译器，建议不使用轮子安装，因为这样可以启用特定架构的优化。

你可以使用以下命令从源代码安装：

pip install sparse_dot_topn --no-binary sparse_dot_topn

构建配置

sparse_dot_topn 在从源代码构建时提供了一些配置选项。从源代码构建可以启用特定架构的优化，建议那些安装了 C++ 编译器的用户使用这种方式。详细信息请参见 INSTALLATION.md。

在集群上分布式处理两个大型（1000万+）稀疏矩阵的前 n 个乘法结果

两个大型（1000万+）稀疏矩阵的前 n 个乘法可以分解成更小的块来处理。例如，可以将稀疏矩阵拆分成只有 100 万行的矩阵，然后对所有这些矩阵对进行（前 n 个）乘法。这样做的原因是为了减少每对矩阵的内存占用，并利用可用的分布式计算能力。

这些对可以分布在集群上计算（例如，我们使用 Spark 集群）。然后将结果矩阵乘积压缩并堆叠，以重现完整的矩阵乘积。

以下是一个示例，展示如何实现这一点，我们将稀疏矩阵 A 中的 1000 行与稀疏矩阵 B 中的 600 行进行匹配， A 和 B 都被分成了块。

import numpy as np
import scipy.sparse as sparse
from sparse_dot_topn import sp_matmul_topn, zip_sp_matmul_topn
# 1a. 示例：在稀疏矩阵A的1000行与稀疏矩阵B的600行之间进行匹配。
A = sparse.random(1000, 2000, density=0.1, format="csr", dtype=np.float32, random_state=rng)
B = sparse.random(600, 2000, density=0.1, format="csr", dtype=np.float32, random_state=rng)

# 1b. 参考完整矩阵乘积的top-n结果
C_ref = sp_matmul_topn(A, B.T, top_n=10, threshold=0.01, sort=True)

# 2a. 拆分稀疏矩阵。这里A被分成五部分，B被分成三部分。
As = [A[i*200:(i+1)*200] for i in range(5)]
Bs = [B[:100], B[100:300], B[300:]]

# 2b. 对所有子矩阵对执行top-n乘法，这里使用双重循环。
# 例如，所有子矩阵对可以分布在集群上并在那里进行乘法运算。
Cs = [[sp_matmul_topn(Aj, Bi.T, top_n=10, threshold=0.01, sort=True) for Bi in Bs] for Aj in As]

# 2c. 对C矩阵进行top-n合并，基于B子矩阵的索引。
Czip = [zip_sp_matmul_topn(top_n=10, C_mats=Cis) for Cis in Cs]

# 2d. 对合并后的C矩阵进行堆叠，基于A子矩阵的索引
# 得到的矩阵C等于C_ref。
C = sparse.vstack(Czip, dtype=np.float32)

迁移到v1版本

sparse_dot_topn v1是对v0.*的重大更改，包含新的绑定和API。新版本增加了对CPython 3.12的支持，现在同时支持整数和浮点数。在内部，我们切换到使用最大堆来收集top-n值，这显著减少了内存占用。之前的实现对top-n选择的复杂度为O(n_columns)，而现在的复杂度为O(top-n)。 awesome_cossim_topn已被弃用，将在未来版本中移除。

用户应该切换到sp_matmul_topn，它在很大程度上兼容：

例如：

C = awesome_cossim_topn(A, B, ntop=10)

可以用以下方式复制：

C = sp_matmul_topn(A, B, top_n=10, threshold=0.0, sort=True)

API变更

ntop已重命名为topn
lower_bound已重命名为threshold
use_threads和n_jobs已合并为n_threads
return_best_ntop选项已移除
test_nnz_max选项已移除
当B的形状不兼容但其转置是兼容时，B会自动转置

return_best_ntop的输出可以通过以下方式复制：

C = sp_matmul_topn(A, B, top_n=10)
best_ntop = np.diff(C.indptr).max()

默认值变更

threshold不再默认为0.0，而是默认禁用

这使得对包含负值的矩阵能够正常工作。此外，在收集非零结果时内部使用了不同的数据结构，内存占用比以前低得多。这意味着threshold参数对性能和内存需求的影响可以忽略不计。如果threshold为None，我们会预先计算非零项的数量，这可以显著减少所需内存，但会有轻微的性能损失（约10%）。