MXNet - 稀疏矩阵的点积

Question

我正在使用 MXNet 构建内容推荐模型。尽管有大约 10K 行，但在 MXNet 中 CPU 和 GPU 上下文会引发内存不足问题。当前代码如下。

```

import mxnet as mx
import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import linear_kernel

df = pd.read_csv("df_text.csv")

tf = TfidfVectorizer(analyzer = "word",
        ngram_range = (1,3),
        min_df = 2,
        stop_words="english")

tfidf_matrix = tf.fit_transform(df["text_column"])

mx_tfidf = mx.nd.array(tfidf_matrix, ctx=mx.gpu())

# Out of memory error occurs here.  
cosine_similarities = mx.ndarray.dot(mx_tfidf, mx_tfidf.T)

```

我知道点积是稀疏矩阵乘以密集矩阵，这可能是问题的一部分。也就是说，是否必须跨多个 GPU 计算点积，以防止出现内存不足问题？

Answer 1

在 MXNet（以及 AFAIK 所有其他平台）中，没有神奇的 "perform dot across GPUs" 解决方案。一种选择是在 MXNet 中使用稀疏矩阵（参见 this tutorial）

另一种选择是通过将输入数组分割成多个矩阵并在每个 GPU 中执行部分点积来实现您自己的多 GPU 点积。