加载Word2Vec向量文件计算余弦相似度时遇MemoryError，求解析

用户问题 #

我加载了一个Word2Vec格式的文件，想要计算向量间的余弦相似度，但不明白出现的问题含义。代码如下：

from gensim.models import Word2Vec 
from sklearn.metrics.pairwise import cosine_similarity 
from gensim.models import KeyedVectors 
import numpy as np 
model = KeyedVectors.load_word2vec_format('it-vectors.100.5.50.w2v') 
similarities = cosine_similarity(model.vectors)

运行后出现MemoryError报错，报错堆栈如下：

--------------------------------------------------------------------------- 
MemoryError Traceback (most recent call last) 
<ipython-input-54-1d4e62f55ebf> in <module>() 
----> 1 similarities = cosine_similarity(model.vectors) 
/usr/local/lib/python3.5/dist-packages/sklearn/metrics/pairwise.py in cosine_similarity(X, Y, dense_output) 
923 Y_normalized = normalize(Y, copy=True) 
924 --> 925 K = safe_sparse_dot(X_normalized, Y_normalized.T, dense_output=dense_output) 
926 
927 return K 
/usr/local/lib/python3.5/dist-packages/sklearn/utils/extmath.py in safe_sparse_dot(a, b, dense_output) 
138 return ret 
139 else: 
--> 140 return np.dot(a, b) 
141 
142 MemoryError: 

请问该错误是什么意思？谢谢！

错误原因与解决办法 #

为什么会出现MemoryError？ #

直白点说就是你的机器内存不够用了！

咱们算笔账就清楚了：从模型文件名it-vectors.100.5.50.w2v能看出，每个词向量的维度是100。当你调用cosine_similarity(model.vectors)时，sklearn会计算所有向量两两之间的相似度，最终生成一个N×N的方阵——N就是模型里的总词数。

举个实际例子：如果模型里有10万个词，这个方阵就有100亿个元素。每个默认的64位浮点数占8字节，算下来这个矩阵要占差不多74.5GB内存！普通机器根本装不下这么大的数据，自然就触发内存报错了。

怎么解决这个问题？ #

根据你的实际需求，给你几个实用方案：

• 只查特定词的相似度（最推荐）：如果你只是想找某个词的相似词，直接用gensim自带的most_similar()方法就好，它会高效计算目标向量和其他向量的相似度，不会生成巨大的矩阵：

  # 示例：查找"computer"最相似的前10个词
  top_similar = model.most_similar("computer", topn=10)
  print(top_similar)
  

• 分批计算相似度：如果确实需要完整的相似度矩阵，可以把向量分成小批次，每次计算一个批次和所有向量的相似度，处理完再合并结果，这样每次只用占用小部分内存：

  batch_size = 1000  # 根据你的内存大小调整批次容量
  all_results = []
  total_vectors = len(model.vectors)
  
  for start_idx in range(0, total_vectors, batch_size):
      end_idx = min(start_idx + batch_size, total_vectors)
      batch_vectors = model.vectors[start_idx:end_idx]
      # 计算当前批次与所有向量的相似度
      batch_similarities = cosine_similarity(batch_vectors, model.vectors)
      all_results.append(batch_similarities)
  
  # 合并所有批次的结果得到完整矩阵
  full_similarity_matrix = np.concatenate(all_results, axis=0)
  

• 压缩向量数据类型：把向量从默认的float64转换成float32，能直接把内存占用减半，对相似度计算的精度影响几乎可以忽略：

  model.vectors = model.vectors.astype(np.float32)
  # 如果还是内存不足，可以结合上面的分批方法一起使用
  

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