将大型 sparknlp 管道加载到 Apache Spark 批处理作业中花费的时间太长

问题描述

我使用 johnsNowlabs 的 SparkNLP 从我的文本数据中提取嵌入，下面是管道。保存到hdfs后模型大小为1.8g

embeddings = BertSentenceEmbeddings.pretrained("labse","xx") \
      .setInputCols("sentence") \
      .setoutputCol("sentence_embeddings")
nlp_pipeline = Pipeline(stages=[document_assembler,sentence_detector,embeddings])
pipeline_model = nlp_pipeline.fit(spark.createDataFrame([[""]]).toDF("text"))

我使用 pipeline_model 将 HDFS 保存到 pipeline_model.save("hdfs:///<path>")。

上面只执行了一次

在另一个脚本中，我正在使用 HDFS 从 pipeline_model = PretrainedPipeline.from_disk("hdfs:///<path>") 加载存储的管道。

上面的代码加载了模型，但是占用的太多了。我在 spark 本地模型（无集群）上对其进行了测试，但我拥有 94g RAM、32 核的高资源。

后来，我在yarn上部署了脚本，其中有 12 个 Executor，每个 Executor 有 3 个内核和 7g ram。我分配了 10g 的驱动程序内存。

脚本再次从 HDFS 加载保存的模型需要太多时间。

当火花到达这一点时（见上面的截图），需要太多时间

我想到了一个方法

预加载

我认为的方法是以某种方式将模型一次预加载到内存中，并且当脚本想要在数据帧上应用转换时，我可以以某种方式调用对预训练管道的引用并在旅途中使用它，而无需执行任何磁盘 I/O。我搜索过，但无处可寻。

请告诉我您对此解决方案的看法以及实现这一目标的最佳方法。

YARN 资源

节点名称	计数	RAM（每个）	核心（每个）
主节点	1	38g	8
次节点	1	38 g	8
工作节点	4	24 g	4
总计	6	172g	32

谢谢

解决方法

正如评论中所讨论的，这是一个基于 PyTorch 的解决方案，而不是 SparkNLP。简化代码：

# labse_spark.py

LABSE_MODEL,LABSE_TOKENIZER = None


def transform(spark,df,input_col='text',output_col='output'):
    spark.sparkContext.addFile('hdfs:///path/to/labse_model')
    output_schema = T.StructType(df.schema.fields + [T.StructField(output_col,T.ArrayType(T.FloatType()))])

    rdd = df.rdd.mapPartitions(_map_partitions_func(input_col,output_col))
    res = spark.createDataFrame(data=rdd,schema=output_schema)
    return res


def _map_partitions_func(input_col,output_col):
    def executor_func(rows):
        # load everything to memory (partitions should be small,~1k rows per partition):
        pandas_df = pd.DataFrame([r.asDict() for r in rows])
        global LABSE_MODEL,LABSE_TOKENIZER
        if not (LABSE_TOKENIZER or LABSE_MODEL):  # should happen once per executor core
            LABSE_TOKENIZER = AutoTokenizer.from_pretrained(SparkFiles.get('labse_model'))
            LABSE_MODEL = AutoModel.from_pretrained(SparkFiles.get('labse_model'))
        
        # copied from HF model card:
        encoded_input = LABSE_TOKENIZER(
            pandas_df[input_col].tolist(),padding=True,truncation=True,max_length=64,return_tensors='pt')
        with torch.no_grad():
            model_output = LABSE_MODEL(**encoded_input)
        embeddings = model_output.pooler_output
        embeddings = torch.nn.functional.normalize(embeddings)

        pandas_df[output_col] = pd.Series(embeddings.tolist())
        return pandas_df.to_dict('records')

    return executor_func