如何在PySpark中加速千万级数据集的多模型推理执行速度

一、完全可以通过向量化替代循环实现大幅提速 #

循环调用10个模型是耗时激增的核心原因——Python循环的开销、逐模型调用的序列化/反序列化成本，加上缺乏批量并行计算，会把单模型的22分钟直接放大数倍。通过向量化改造，能把多模型推理的计算逻辑整合到一次批量处理中：

• 将10个TensorFlow模型合并为一个联合推理图：用TensorFlow的Functional API共享输入层，把所有模型的输出拼接后直接计算均值。示例代码如下：

  import tensorflow as tf
  from tensorflow.keras.models import Model
  
  # 加载10个预训练模型
  models = [tf.keras.models.load_model(f"model_{i}.h5") for i in range(10)]
  
  # 共享输入层（假设所有模型输入格式一致）
  input_layer = models[0].input
  # 获取所有模型的输出
  outputs = [model(input_layer) for model in models]
  # 计算均值作为最终输出
  avg_output = tf.keras.layers.average(outputs)
  
  # 创建联合模型
  ensemble_model = Model(inputs=input_layer, outputs=avg_output)
  

  改造后只需调用一次ensemble_model.predict()，就能批量处理数据并直接得到均值得分，完全消除循环开销，同时TensorFlow会自动对整个计算图做算子融合、并行计算等优化。

二、其他关键加速方法 #

1. 用Pandas UDF替代普通Python UDF

PySpark普通Python UDF是逐行处理，性能极差。改用Vectorized Pandas UDF（基于Apache Arrow），可批量处理整个分区的Pandas数据，大幅减少Python与JVM之间的数据传输开销：

from pyspark.sql.functions import pandas_udf
import pandas as pd

# 提前加载联合模型，可通过广播变量分发到executor避免重复加载
ensemble_model = ...

@pandas_udf("double")
def predict_udf(input_features: pd.Series) -> pd.Series:
    # 转换为模型输入格式
    X = input_features.to_numpy().reshape(-1, feature_dim)
    # 批量推理，设置合适的batch_size
    predictions = ensemble_model.predict(X, batch_size=2048)
    return pd.Series(predictions.flatten())

# 应用UDF到目标DataFrame
result_df = df.withColumn("avg_score", predict_udf("features"))

2. 优化Spark集群资源与分区

• 调整分区数：确保每个分区大小在100MB-200MB之间。针对8000万-1亿行的数据集，若每行约1KB，总数据量约800GB，建议设置8000-10000个分区，避免分区过大导致内存溢出、过小带来过多调度开销。
• 分配充足硬件资源：每个executor分配4-8核CPU、16-32GB内存；若有GPU资源，配置Spark GPU调度（如利用RAPIDS或TensorFlow on Spark），让模型推理在GPU上并行执行，批量计算能力可将速度提升数倍。

3. 批量推理参数调优

• 调整batch_size：GPU环境可设为1024-4096，CPU环境设为512-1024，让TensorFlow充分利用硬件计算核心。
• 关闭Eager Execution：启用Graph模式避免即时执行开销：

  tf.compat.v1.disable_eager_execution()
  

4. 数据预处理优化

• 优先用Spark内置算子完成特征工程（如VectorAssembler、StandardScaler），避免Python UDF——Spark内置算子基于JVM实现，性能远高于Python UDF。
• 用Parquet格式存储数据集，列式存储的高压缩率和快速读取能力可减少IO耗时。

5. 模型轻量化优化

• 模型量化：用TensorFlow的tf.quantization工具将模型转为INT8量化模型，减少体积和计算量，推理速度提升2-3倍且精度损失极小。
• 模型蒸馏：若业务允许，将10个模型的知识蒸馏到一个更小的模型中，既能保留集成精度，又能大幅降低推理耗时。

内容的提问来源于stack exchange，提问作者krishna kaushik