大向量余弦相似度计算遇内存错误，求批量处理或算法优化方案

问题根源 #

计算(94955,94955)规模的余弦相似度矩阵需要约67GiB内存，远超普通机器的内存上限，这是导致MemoryError的核心原因——不管用稀疏还是稠密存储，全量保存这个矩阵都不现实。

方案一：批量处理+保留Top-K结果 #

不需要一次性计算所有样本的相似度，只针对需求计算并保留每个样本的Top-N相似结果，避免存储完整的相似度矩阵。

代码示例（分批次计算Top-K）

from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# 设定每个样本需要保留的相似样本数量
top_k = 10
similarity_results = []
batch_size = 1000  # 每次处理的样本批次大小，可根据内存调整
n_samples = vectors.shape[0]

for i in range(0, n_samples, batch_size):
    # 取出当前批次的向量
    current_batch = vectors[i:i+batch_size]
    # 计算当前批次与所有样本的余弦相似度（返回稀疏矩阵）
    batch_sim = cosine_similarity(current_batch, vectors, dense_output=False)
    
    # 遍历批次内的每个样本，提取Top-K相似项
    for row in batch_sim:
        # 稀疏矩阵转数组，排序后取Top-K（排除样本自身）
        sim_values = row.toarray().flatten()
        sorted_indices = np.argsort(sim_values)[::-1][1:top_k+1]
        sorted_scores = sim_values[sorted_indices]
        similarity_results.append((sorted_indices, sorted_scores))

如果是针对单个样本做推荐，直接单独计算该样本与所有其他样本的相似度即可，无需全量处理：

# 假设target_idx是目标样本的索引
target_vector = vectors[target_idx].reshape(1, -1)
sim_scores = cosine_similarity(target_vector, vectors)[0]
# 取Top-K相似样本
top_indices = np.argsort(sim_scores)[::-1][1:top_k+1]

方案二：更换为近似最近邻（ANN）算法 #

如果业务可以接受一定的近似误差，使用专门的大规模相似度检索算法是更高效的选择，这类算法能大幅降低内存占用和计算时间。

示例：使用FAISS（Facebook开源）

FAISS针对高维向量的近邻检索做了深度优化，适合百万级以上样本：

import faiss
import numpy as np

# 将向量转换为FAISS要求的float32格式
vectors_float32 = vectors.astype(np.float32)
# 归一化向量（余弦相似度等价于归一化后的L2距离的补）
faiss.normalize_L2(vectors_float32)

# 构建IVF_FLAT索引（适合大规模数据）
n_clusters = 100  # 聚类中心数量，可根据样本量调整
quantizer = faiss.IndexFlatL2(5000)  # 基础量化器
index = faiss.IndexIVFFlat(quantizer, 5000, n_clusters, faiss.METRIC_L2)

# 训练索引并添加向量
index.train(vectors_float32)
index.add(vectors_float32)

# 检索每个样本的Top 10相似项
top_k = 10
distances, indices = index.search(vectors_float32, top_k)
# 转换为余弦相似度值（L2距离越小，余弦相似度越高）
cosine_similarities = 1 - distances

其他可选的ANN库还有Annoy、HNSWlib，用法类似，都能有效解决大规模向量的相似度检索问题。

方案三：结合降维（辅助优化） #

可以先用PCA将5000维的向量降到更低维度（如200维），再结合上述批量处理或ANN算法，进一步降低计算和内存压力，但会损失少量精度：

from sklearn.decomposition import PCA

# 降维到200维
pca = PCA(n_components=200)
vectors_reduced = pca.fit_transform(vectors)

# 之后可以用批量处理或FAISS处理降维后的向量

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