当我们训练好了一个模型,想在大规模的数据上预测时,pyspark的Pandas UDF在加载模型并使用模型对数据进行分布式预测很有用。
但在spark 3.0 以前,当模型很大时,因为在同一个python工作进程中,每个批次Pandas UDF都会反复加载同一个模型,会造成很高的额外开销。
Spark 3.0 以后, 因为Pandas UDF支持Series to Series, Iterator of Series to Iterator of Series 等,所以就只用加载一次模型。
下面的示例代码给出了一个使用pyspark 调用模型训练的示例
# 一种模型预测示例
import os
from pyspark.sql import functions as F
from pyspark.sql.functions import pandas_udf
# 这里的定义变量是因为模型存储在s3上
filename = '/tmp/catboost_model_1'
dest_path = 'remote_filesystem/catboost_model_1'
@pandas_udf("float")
def predict(features):
# 定义模型
model_cb = CatBoostClassifier()
# 如果不能保证spark的每个节点上都已经有模型了,一定要作下面的判断,不然会报错类似 Model file doesn't exist: /tmp/catboost_model_1 的错误
if not os.path.exists(filename):
cmd = '/usr/lib/hadoop/bin/hadoop' + '-get -f {} {}'.format(dest_path, filename)
os.system(cmd)
# 加载模型
model_cb.load_model(filename, format='cbm')
logger.debug('finished load model')
# 预测
# return pd.Series(model_cb.predict(features)) # 直接预测结果
return pd.Series(model_cb.predict_proba(features)[:, 1]) # 预测概率
# 将输入特征列组成结构体
feature_col = F.struct(F.col('feature1'), F.col('feature2'), F.col('feature3'))
df = df.withColumn('prediction', predict(feature_col))