如何使用 Spark 进行大型模型的机器学习工作流

Posted 2023-03-31

【中文标题】如何使用 Spark 进行大型模型的机器学习工作流

【英文标题】：How to use Spark for machine learning workflows with big models

【发布时间】：2020-07-01 11:59:32

【问题描述】：

有没有一种内存高效的方法可以将大型 (>4GB) 模型应用到 Spark Dataframes 而不会遇到内存问题？

我们最近将自定义管道框架移植到 Spark（使用 python 和 pyspark），但在将 Word2Vec 和自动编码器等大型模型应用于标记化文本输入时遇到了问题。首先，我非常天真地将转换调用转换为udfs（pandas 和 spark “native” 的），这很好，只要使用的模型/实用程序足够小，可以广播或重复实例化：

@pandas_udf("array<string>")
def tokenize_sentence(sentences: pandas.Series):
    return sentences.map(lambda sentence: tokenize.word_tokenize(sentence))

对大型模型尝试相同的方法（例如，通过 word2vec 将这些标记嵌入向量空间）导致性能很差，我明白了原因：

@pandas_udf("array<array<double>>")
def rows_to_lists_of_vectors(rows):
    model = api.load('word2vec-google-news-300')

    def words_to_vectors(words) -> List[List[float]]:
        vectors = []
        for word in words:
            if word in model:
                vec = model[word]
                vectors.append(vec.tolist())
        return vectors
    return rows.map(words_to_vectors)

上面的代码会重复实例化 ~4Gb word2vec 模型，将其从磁盘加载到 RAM，这非常慢。我可以通过使用mapPartition 来解决这个问题，它至少每个分区只加载一次。但更重要的是，如果我没有严格限制任务的数量，这会导致与内存相关的问题（至少在我的开发机器上）崩溃，这反过来又会使小型 udf 变得非常慢。例如，将任务数限制为 2 可以解决内存崩溃问题，但会使标记化变得非常缓慢。

我知道 Spark 中有一个完整的管道框架，可以满足我们的需求，但在提交之前，我想了解我遇到的问题是如何在那里解决的。也许我们可以使用一些关键实践，而不必重写我们的框架。

因此，我的实际问题是双重的：

假设我们为步骤编写了自定义的 Estimator 和 Transformer，Spark 开箱即用（例如 Tokenizer 和 Word2Vec）未涵盖这些步骤，那么使用 Spark 管道框架是否可以解决我们关于性能和内存的问题。 Spark 如何解决这些问题（如果有的话）？我可以改进当前的方法吗，或者使用 python 是不可能的（据我了解，进程不共享内存空间）。

如果以上任何内容让您认为我错过了 Spark 的核心原则，请指出，毕竟我刚刚开始使用 Spark。

【参考方案1】：

这取决于各种因素（模型、集群资源、管道），但试图回答您的主要问题：

1)。如果 Spark 管道在 Tokenizers、Words2Vec 等方面满足您的需求，它们可能会解决您的问题。但是，这些管道不如现成的可用并加载了 api.load 的那么强大。你可能还想看看 Deeplearning4J，它将这些带到 Java/Apache Spark 中，看看它如何做同样的事情：tokenize、word2vec 等

2)。按照当前的方法，我会看到在foreachParition 或mapPartition 中加载模型，并确保模型可以适合每个分区的内存。您可以根据集群资源将分区大小缩小到更实惠的数量以避免内存问题（例如，当您为每个分区创建一个数据库连接而不是为每一行创建一个数据库连接时也是如此）。

通常，当您应用一种对 Spark 友好且不混入第三个外部方的业务逻辑时，Spark udfs 会很好。