计算Tfidf矩阵和预测向量之间的相似性会导致内存溢出

Question

我使用以下代码在〜20,000,000个文档上生成了tf-idf模型，该模型效果很好。问题是，当我尝试使用计算时的相似性指标linear_kernel内存使用情况炸毁：

from sklearn.feature_extraction.text import TfidfVectorizer 
from sklearn.metrics.pairwise import linear_kernel 

train_file = "docs.txt" 
train_docs = DocReader(train_file) #DocReader is a generator for individual documents 

vectorizer = TfidfVectorizer(stop_words='english',max_df=0.2,min_df=5) 
X = vectorizer.fit_transform(train_docs) 

#predicting a new vector, this works well when I check the predictions 
indoc = "This is an example of a new doc to be predicted" 
invec = vectorizer.transform([indoc]) 

#This is where the memory blows up 
similarities = linear_kernel(invec, X).flatten() 

好像这不应该占用多少内存，做一个1行-CSR的到20密尔的比较-row-CSR应输出一个1x20mil的ndarray。 Justy FYI：X是一个CSR矩阵〜12 GB的内存（我的电脑只有16个）。我试图寻找gensim来取代这一点，但我找不到一个很好的例子。

关于我失踪的任何想法？

Answer 1

您可以分批处理。这里是一个基于你的代码片段的例子，但将数据集替换为sklearn中的某个东西。对于这个较小的数据集，我计算它的原始方式以及显示结果是相同的。你可以使用更大的批量。

import numpy as np 
from sklearn.feature_extraction.text import TfidfVectorizer 
from sklearn.metrics.pairwise import linear_kernel 
from sklearn.datasets import fetch_20newsgroups 

train_docs = fetch_20newsgroups(subset='train') 

vectorizer = TfidfVectorizer(stop_words='english', max_df=0.2,min_df=5) 
X = vectorizer.fit_transform(train_docs.data) 

#predicting a new vector, this works well when I check the predictions 
indoc = "This is an example of a new doc to be predicted" 
invec = vectorizer.transform([indoc]) 

#This is where the memory blows up 
batchsize = 1024 
similarities = [] 
for i in range(0, X.shape[0], batchsize): 
    similarities.extend(linear_kernel(invec, X[i:min(i+batchsize, X.shape[0])]).flatten()) 
similarities = np.array(similarities) 
similarities_orig = linear_kernel(invec, X) 
print((similarities == similarities_orig).all()) 

输出：

True