Got memory error while performing cosine similarity on a huge vector

Viewed 51

I was trying to build a content based recommendation system using bag of words model. The tutorial which I am following uses cosine similarity from sklearn library on a vector of size (4000,5000) where 4000 is the number of rows in dataset and 5000 is the number of features.

from sklearn.feature_extraction.text import CountVectorizer
cv = CountVectorizer(max_features=5000, stop_words='english')
vectors = cv.fit_transform(new_df['tags']).toarray() 
// here new_df is the dataframe and new_df[tags] contain all the tags (eg: location, genre) based on which recommendation will be performed

But when I try to implement cosine similarity on another datset with 94955 rows, which results in a vector of size (94955, 5000), I get the following error

MemoryError: Unable to allocate 67.2 GiB for an array with shape (94955, 94955) and data type float64

on line similarity = cosine_similarity(vectors, dense_output=False)

Is there a way to implement batching on cosine similarity so that I can overcome this issue or should I change the algorithm?

1 Answers

sklearn.metrics.pairwise.cosine_similarity() computes pairwise similarities between all samples in vectors and returns array of shape (94955, 94955). Each computation is performed on 5000-dimensional vector. That is too memory costlier for a single machine, and possibly can't scale vertically.

This problem can be solved by scaling horizontally. You can use spark for this.

We had a similar requirement to compute pairwise similarity on huge data. I wrote a vectorised cosine similarity computation code for this. It took few hours; but finally completed, in my case. There were 75,000 x 75,000 similarity comparisons with 1000-dimensional vectors.

The idea is to transform the dataframe in large number of (a,b) pairs, and apply following vectorised operations on columns. Spark will distribute pairs on different worker machines.

First, convert features into two vector columns using pyspark.ml.feature.VectorAssembler. The two vectors will be used to compute pairwise similarities.

from pyspark.ml.feature import VectorAssembler
feature_col_names = [...] # TODO - collect your feature column names
assembler = VectorAssembler(inputCols=feature_col_names, outputCol="features_vec_1")
df = assembler.transform(df)

from pyspark.sql.functions import col
df = df.withColumn("features_vec_2", col("features_vec_1"))

Compute cross-join of the two vector columns for pairwise similarities:

df = df.select("features_vec_1").crossJoin(df.select("features_vec_2"))

Apply following vectorised column operations to compute cosine similarity:

import pyspark.sql.functions as fn
df = df.withColumn("axb", fn.zip_with("features_vec_1", "features_vec_2", lambda x, y: x * y))
df = df.withColumn("axa", fn.zip_with("features_vec_1", "features_vec_1", lambda x, y: x * y))
df = df.withColumn("bxb", fn.zip_with("features_vec_2", "features_vec_2", lambda x, y: x * y))
df = df.withColumn("dot_prod", fn.expr("AGGREGATE(axb, DOUBLE(0), (a,x)->a+x)"))
df = df.withColumn("norm_1", fn.expr("sqrt(AGGREGATE(axa, DOUBLE(0), (a,x)->a+x))"))
df = df.withColumn("norm_2", fn.expr("sqrt(AGGREGATE(bxb, DOUBLE(0), (a,x)->a+x))"))
df = df.withColumn("cosine_similarity", fn.lit(fn.col("dot_prod") / (fn.col("norm_1") * fn.col("norm_2"))))
df = df.drop("axb", "axa", "bxb", "dot_prod", "norm_1", "norm_2")

You may want to apply some termination operation (like collect() or df.write) to trigger the above computations.

df.write.mode("overwrite").parquet("/file/path")
df = sparkSession.read.parquet("/file/path")

OR

similarity_list = df.collect()

At this point you've cosine similarities computed in following pairwise fashion:

+--------------+--------------+-----------------+
|features_row_1|features_row_2|cosine_similarity|
+--------------+--------------+-----------------+
|         row_1|         row_1|              1.0|
|         row_1|         row_2|              0.5|
|         row_1|         ...  |                 |
|         row_1|     row_94955|              0.7|
|         row_2|         row_1|              0.5|
|         row_2|         row_2|              1.0|
|         row_2|         ...  |                 |
|         row_2|     row_94955|              0.2|
|     row_94955|         row_1|              0.7|
|     row_94955|         row_2|              0.2|
|     row_94955|         ...  |                 |
|     row_94955|     row_94955|              1.0|
+--------------+--------------+-----------------+
Related