Jupyter Notebook大型矩阵余弦相似度计算内存错误解决方案咨询

针对你8GB内存的设备，以下是几个可行的高效计算方法，避免内存错误：

1. 利用稀疏矩阵压缩存储 #

计数矩阵通常存在大量零值，转成稀疏矩阵后仅存储非零元素，能大幅降低内存占用，sklearn.metrics.pairwise.cosine_similarity原生支持稀疏矩阵输入：

from scipy.sparse import csr_matrix
from sklearn.metrics.pairwise import cosine_similarity

# 假设你的原始计数矩阵为count_matrix（numpy数组）
sparse_count_matrix = csr_matrix(count_matrix)
# 计算得到的相似度矩阵也是稀疏格式
cos_sim_sparse = cosine_similarity(sparse_count_matrix)

这种方式能把原本几十GB的内存需求降到几百MB级别，完全适配你的设备。

2. 分块计算+磁盘持久化 #

如果稀疏矩阵仍有压力，可将矩阵拆分成小块分批计算，每块计算完成后立即保存到磁盘，避免一次性占用大量内存：

import numpy as np

# 每次处理的行数，可根据内存调整（比如500或1000）
block_size = 1000
total_rows = count_matrix.shape[0]

for start_idx in range(0, total_rows, block_size):
    end_idx = min(start_idx + block_size, total_rows)
    # 计算当前块与全量矩阵的余弦相似度
    block_sim = cosine_similarity(count_matrix[start_idx:end_idx], count_matrix)
    # 保存到磁盘，避免驻留内存
    np.savez(f"cos_sim_block_{start_idx}_{end_idx}.npz", sim_block=block_sim)

后续需要使用时，可按需加载单个分块文件。

3. 计算Top-K近似相似度（推荐系统专属优化） #

推荐系统通常只需要每个用户/物品的Top-K相似项，无需全量相似度矩阵，可使用近似最近邻算法减少计算量和内存占用：

from sklearn.neighbors import NearestNeighbors
from scipy.sparse import csr_matrix

sparse_matrix = csr_matrix(count_matrix)
# 初始化模型，设置需要获取的相似项数量（比如Top20）
nn_model = NearestNeighbors(n_neighbors=20, metric='cosine', algorithm='brute')
nn_model.fit(sparse_matrix)

# 获取每个样本的Top20相似项索引和余弦距离
distances, top_k_indices = nn_model.kneighbors(sparse_matrix)
# 余弦相似度 = 1 - 余弦距离
top_k_similarities = 1 - distances

这种方式的内存占用仅与样本数和Top-K数量相关，8GB内存完全够用。

4. 降低数据类型精度 #

将原始矩阵的int64类型转换为更低精度的数值类型，直接减少内存占用：

import numpy as np

# 转换为float32，内存占用减半（从8字节/元素降到4字节/元素）
count_matrix_float32 = count_matrix.astype(np.float32)
cos_sim = cosine_similarity(count_matrix_float32)

# 如果计数矩阵的数值范围较小（比如小于255），可尝试uint8（1字节/元素）
# count_matrix_uint8 = count_matrix.astype(np.uint8)

注意转换前需确认数值不会溢出，比如uint8仅支持0-255的数值。

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