如何在PySpark DataFrame上执行自定义PyTorch模型的推理任务？

错误根源 #

报错核心是Spark与PyTorch的数据格式不匹配：

• 用concat构造的features列是多嵌套数组的拼接，不符合模型预期的6个独立输入张量格式；
• 冗余的数组嵌套操作（如[x]把一维数组转成二维）导致数据维度混乱，Spark传递给模型的列表无法被正确转换为张量；
• 未适配Spark到PyTorch的张量类型转换，且缺少推理时的分数计算逻辑。

修复步骤与代码调整 #

1. 清理预处理代码，移除冗余操作 #

删除多余的UDF嵌套，确保每个输入列是一维整数数组（Spark会自动将batch内的样本堆叠为二维张量）：

data=[("[CLS] My name is [unused0] John [unused1] ! [SEP]","[CLS] John [unused3] male given name [SEP]")]
spark = sparknlp.start()
df=spark.createDataFrame(data).toDF("context","candidate")

# 预处理：生成一维数组格式的输入特征，无需额外嵌套
df = df.withColumn("ctx_idx", udf(lambda x: token_to_id(x,32), ArrayType(IntegerType()))("context"))
df = df.withColumn("ctx_segment", udf(lambda x: [0]*len(x), ArrayType(IntegerType()))("ctx_idx"))
df = df.withColumn("ctx_mask", udf(lambda x: [1 if i !=0 else 0 for i in x], ArrayType(IntegerType()))("ctx_idx"))

df = df.withColumn("cnd_idx", udf(lambda x: token_to_id(x,512), ArrayType(IntegerType()))("candidate"))
df = df.withColumn("cnd_segment", udf(lambda x: [0]*len(x), ArrayType(IntegerType()))("cnd_idx"))
df = df.withColumn("cnd_mask", udf(lambda x: [1 if i !=0 else 0 for i in x], ArrayType(IntegerType()))("cnd_idx"))

# 删除未定义的run_model UDF代码行
# df =df.withColumn("score", udf(lambda ctx_ids, cnd_ids : run_model(ctx_ids, cnd_ids), ArrayType(IntegerType()))("context","candidate"))

2. 封装推理专用模型类 #

原模型仅返回embedding，封装适配Spark的推理类，完成张量类型转换、embedding分数计算，并确保模型处于评估模式：

class BiEncoderInferenceModule(torch.nn.Module):
    def __init__(self, original_model):
        super().__init__()
        self.original_model = original_model
        # 切换为评估模式，禁用训练相关层（如Dropout）
        self.original_model.eval()

    def forward(self, ctx_ids, ctx_seg, ctx_mask, cnd_ids, cnd_seg, cnd_mask):
        # 将Spark传入的张量转换为Bert所需的Long类型
        ctx_ids = ctx_ids.long()
        ctx_seg = ctx_seg.long()
        ctx_mask = ctx_mask.long()
        cnd_ids = cnd_ids.long()
        cnd_seg = cnd_seg.long()
        cnd_mask = cnd_mask.long()

        # 调用原模型获取embedding
        embedding_ctxt, embedding_cands = self.original_model(
            ctx_ids, ctx_seg, ctx_mask, cnd_ids, cnd_seg, cnd_mask
        )

        # 计算匹配分数，与本地推理逻辑一致
        embedding_ctxt = embedding_ctxt.unsqueeze(1)
        embedding_cands = embedding_cands.unsqueeze(2)
        scores = torch.bmm(embedding_ctxt, embedding_cands)
        scores = torch.squeeze(scores)
        return scores

3. 创建并运行Spark Pipeline #

使用create_spark_torch_model时，直接指定6个输入列，而非拼接后的单一列，确保模型能正确接收每个输入张量：

# 包装原模型为推理专用模型
inference_model = BiEncoderInferenceModule(torch_model)

# 创建Spark兼容的Torch模型，指定多输入列
spark_torch_model = create_spark_torch_model(
    inference_model,
    inputCols=["ctx_idx", "ctx_segment", "ctx_mask", "cnd_idx", "cnd_segment", "cnd_mask"],
    predictionCol='score'
)

# 构建Pipeline并执行推理
pipeline = Pipeline(stages=[spark_torch_model])
result = pipeline.fit(df).transform(df)

# 查看结果
result.select("context", "candidate", "score").show()

额外注意事项 #

• 确保Spark集群的每个节点都安装了PyTorch、Transformers等依赖库；
• 若模型较大，建议广播模型到每个节点，避免重复加载；
• 如需使用GPU，需配置Spark的GPU支持参数（如spark.task.resource.gpu.amount）。

内容的提问来源于stack exchange，提问作者Hamed Parvaresh