如何高效在CSR格式TF-IDF向量上计算训练与测试数据的相似度

针对你2200万条训练数据的TF-IDF CSR矩阵，直接转稠密数组崩溃完全是预期结果——高维稀疏向量转稠密后内存占用会呈指数级增长，必须全程基于稀疏矩阵操作来实现高效相似度计算，以下是具体方案：

核心原则：全程保留稀疏格式 #

绝对避免使用toarray()或todense()转换训练矩阵，所有运算都基于CSR（或CSC）稀疏矩阵执行，这是控制内存的关键。

推荐方案1：基于K近邻快速获取Top N相似样本 #

如果你的需求是找到每个测试样本的Top K最相似训练样本（而非计算所有样本的相似度），这是最高效的方式：

1. 归一化TF-IDF向量
   余弦相似度等价于归一化向量的点积，先对训练和测试的TF-IDF矩阵做L2归一化，确保后续计算的点积就是余弦相似度：

   from sklearn.preprocessing import normalize
   
   # 训练矩阵归一化，结果仍为CSR格式
   train_csr_norm = normalize(train_csr, norm='l2', axis=1)
   # 测试矩阵同样归一化
   test_csr_norm = normalize(test_csr, norm='l2', axis=1)
   

2. 用暴力搜索K近邻模型计算
   高维稀疏向量不适合KD树等空间索引算法，直接用暴力搜索配合多线程加速：

   from sklearn.neighbors import NearestNeighbors
   
   # 初始化模型，设置需要返回的近邻数、余弦度量、多线程
   nn_model = NearestNeighbors(
       n_neighbors=10,  # 按需修改需要的Top K数量
       metric='cosine',
       algorithm='brute',
       n_jobs=-1  # 用所有CPU核心加速
   )
   nn_model.fit(train_csr_norm)
   
   # 获取每个测试样本的Top K近邻：distances是余弦距离（1-余弦相似度），indices是训练样本的索引
   distances, indices = nn_model.kneighbors(test_csr_norm)
   # 转换为余弦相似度
   similarities = 1 - distances
   

   这种方式只会存储每个测试样本的K个相似度值，内存占用极低，计算速度也远快于全量相似度计算。

推荐方案2：全量相似度计算（仅当必须时使用） #

如果确实需要计算每个测试样本与所有训练样本的相似度，依然用稀疏矩阵的点积操作：

# 归一化后的测试矩阵点乘训练矩阵的转置，得到相似度矩阵（CSR格式）
similarity_matrix = test_csr_norm.dot(train_csr_norm.T)

结果是一个形状为[测试样本数, 2200万]的CSR矩阵，仅存储非零相似度值，内存占用远低于稠密数组。你可以通过similarity_matrix[i].data和similarity_matrix[i].indices获取第i个测试样本的非零相似度值及对应的训练样本索引。

内存优化进阶技巧 #

• 分块处理训练数据：如果训练矩阵大到单块内存放不下，可以将其拆分为多个子块，依次计算测试数据与每个子块的相似度，最后合并结果（比如找Top K时，每块保留Top K，再全局排序）。
• 使用Faiss加速稀疏向量搜索：Facebook的Faiss库对稀疏向量的余弦相似度搜索做了专门优化，性能比sklearn的暴力搜索更优，适合超大规模样本场景。

内容的提问来源于stack exchange，提问作者Vikram