如何缩小Numpy数据类型？解决cosine_similarity内存不足问题

方法一：直接指定计算时的数据类型 #

sklearn的cosine_similarity函数支持通过dtype参数指定输出矩阵的数据类型，无需先生成内存占用大的float64矩阵再转换。将类型改为np.float32可直接把内存占用减半：

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

similarity_matrix = cosine_similarity(tfidf_matrix, dtype=np.float32)

float32的精度足够覆盖大多数余弦相似度计算场景，完全能满足你的需求。

方法二：保留两位小数进一步压缩内存 #

如果确实只需要保留两位小数，可以在计算后对结果做四舍五入，再转换为更紧凑的数据类型：

1. 使用float16类型（内存仅为float64的1/4）：

similarity_matrix = cosine_similarity(tfidf_matrix, dtype=np.float32)
similarity_matrix = np.round(similarity_matrix, 2).astype(np.float16)

2. 转换为整数类型（内存占用最小）：
   将相似度值乘以100后转成int16，后续使用时再除以100即可：

similarity_matrix = cosine_similarity(tfidf_matrix, dtype=np.float32)
similarity_matrix = np.round(similarity_matrix * 100).astype(np.int16)
# 后续需要使用时还原：
# similarity_matrix = similarity_matrix / 100.0

注意事项 #

• float16的精度有限，极端情况下可能出现精度损失，但对于保留两位小数的需求完全足够。
• 优先使用dtype参数直接指定计算类型，比先生成float64再转换更节省内存和计算时间。

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