Le calcul de la similarité entre la matrice Tfidf et le vecteur prédit provoque un débordement de la mémoire

Question

J'ai généré un modèle tf-idf sur ~ 20 000 000 de documents en utilisant le code suivant, qui fonctionne bien. Le problème est quand j'essaie de calculer les scores de similarité lors de l'utilisation linear_kernel l'utilisation de la mémoire explose:

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() 

On dirait que cela ne devrait pas prendre beaucoup de mémoire, en faisant une comparaison d'une 1-ligne-RSE à un 20mil -row-CSR devrait sortir un ndarray 1x20mil.

Justy FYI: X est une matrice CSR ~ 12 Go en mémoire (mon ordinateur n'en a que 16). J'ai essayé de remplacer ceci, mais je ne trouve pas un bon exemple.

Des pensées sur ce que je manque?

Answer 1

Vous pouvez effectuer le traitement par lots. Voici un exemple basé sur votre fragment de code mais en remplaçant l'ensemble de données par quelque chose dans sklearn. Pour ce jeu de données plus petit, je le calcule également de façon originale pour montrer que les résultats sont équivalents. Vous pouvez probablement utiliser une taille de lot plus grande.

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()) 

Sortie:

True