計算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